hi ptrblck, thanks for the reply. Sorry, what do you mean by “complete stack trace”? The full error is THCudaCheck FAIL file=/pytorch/aten/src/THC/THCGeneral.cpp line=383 error=11 : invalid argument. It somehow doesn’t run after throwing this error.
I have the code of
if torch.cuda.device_count() > 1: 
    model = torch.nn.DataParallel(model)

to “put” my model on the 2 GPUs. Other than that, it works fine on my CPU.
edit: it works on accelerator="type=nvidia-tesla-p4,count=2", i.e., no error thrown. Apparently it only fails for t4.