Multiprocessing in OpenAI Gym with abseil

Question

I am struggling with multiprocessing in OpenAI Gym with the abseil library. Basically, the gym.make seems working. However, I am trying to use gym-super-mario-bros which is not working. Below is a minimal working example:

from absl import app
import os
os.environ['OMP_NUM_THREADS'] = '1'

import gym
import gym_super_mario_bros
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT
from nes_py.wrappers import JoypadSpace
import multiprocessing as mp
import torch
from torch import nn
import time


def get_env():
    env = JoypadSpace(gym_super_mario_bros.make('SuperMarioBros-1-1-v0'), SIMPLE_MOVEMENT)
    # env = gym.make('LunarLander-v2') # other environment such as this one and others works well
    return env

def do_something(env, net1, net2):
    print('inside do_something')
    obs = env.reset()
    print(f'after reset {obs.shape}')
    net2.load_state_dict(net1.state_dict())
    print('after load_state_dict')

def main(args):
    del args
    env = get_env()

    net1 = nn.Sequential(nn.Conv2d(1, 20, 5), nn.ReLU())
    net2 = nn.Sequential(nn.Conv2d(1, 20, 5), nn.ReLU())

    net1.share_memory()
    net2.share_memory()

    device = torch.device('cuda')
    net1 = net1.to(device)
    net2 = net2.to(device)

    p = mp.Process(target=do_something, args=(env, net1, net2,))
    p.start()

    time.sleep(4.0) # wait for the above process to execute print statements
    env.close()

if __name__ == '__main__':
    mp.set_start_method('spawn')
    app.run(main)

Upon running the above code, it seems stuck after printing inside do_something. It also reports CUDA warning in the terminal shown below:

$ python mwe.py 
inside do_something
terminate called after throwing an instance of 'std::bad_alloc'
  what():  std::bad_alloc
[W CudaIPCTypes.cpp:15] Producer process has been terminated before all shared CUDA tensors released. See Note [Sharing CUDA tensors]
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)
[W CUDAGuardImpl.h:46] Warning: CUDA warning: driver shutting down (function uncheckedGetDevice)
[W CUDAGuardImpl.h:62] Warning: CUDA warning: invalid device ordinal (function uncheckedSetDevice)

Please note that it works smoothly without any CUDA warning in case of the gym.make.

Version Info.

I am using Ubuntu 22.04.2 LTS OS. Below is the version info:

Library	Version
absl-py	1.3.0
cuda	11.7
gym	0.17.2
gym-super-mario-bros	7.4.0
nes-py	8.2.1
numpy	1.21.0
python	3.9.16
torch	1.13.1

Is there a workaround to make mario work in multiprocessing environment?