Hi
I think you are right, and it is indeed a red herring. The root cause might be related to hardware. The following code snippet can reproduce the bug:
>>> import torch
>>> torch.rand(1).cuda(0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: CUDA error: out of memory
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
>>> torch.rand(1).cuda(1)
tensor([0.7275], device='cuda:1')

It seems that there are some problems with GPU 0. Do you have any idea?