I’m able to run the code until the numpy error is raised:
tensor(14018246., device='cuda:0')
torch.Size([2, 3, 32, 180, 320])
max: tensor(1.9529, device='cuda:0')
max: tensor(0.9529, device='cuda:0')
tensor(13985083., device='cuda:0')
torch.Size([2, 3, 32, 180, 320])
max: tensor(1.9490, device='cuda:0')
max: tensor(0.9490, device='cuda:0')
input to fg_mpi:  tensor(113580.9062, device='cuda:0')
Traceback (most recent call last):
  File "reproduce_bug.py", line 194, in <module>
    net(sample)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 881, in _call_impl
    result = self.forward(*input, **kwargs)
  File "reproduce_bug.py", line 78, in forward
    save_rgba(fg_mpi.squeeze(), 'rgba')
  File "/workspace/src/visualization.py", line 19, in save_rgba
    plt.imsave(os.path.join(folder, f'rgba_{i:03d}.png'), tmp.cpu().numpy())
  File "/opt/conda/lib/python3.8/site-packages/matplotlib/pyplot.py", line 2251, in imsave
    return matplotlib.image.imsave(fname, arr, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/matplotlib/image.py", line 1581, in imsave
    image = PIL.Image.frombuffer(
  File "/opt/conda/lib/python3.8/site-packages/PIL/Image.py", line 2637, in frombuffer
    im = im._new(core.map_buffer(data, size, decoder_name, None, 0, args))
ValueError: ndarray is not C-contiguous

Could you let me know, which tensors are supposed to show the potential race condition for further debugging?