RosettaFold2 Colab crashing at end of prediction run

At the end of a RosettaFold2 Colab dimer prediction run (with default variables) that has run through 6 recycles––where I tried both V100 GPU & then T4 TPU settings (both with High Ram option)––the program crashes with the following error message below.

The RuntimeError: "Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_bmm)**

jobname: XXX
lengths: [555, 555]
getting unpaired MSA
COMPLETE: 100%|██████████| 150/150 [elapsed: 00:00 remaining: 00:00]
N=2103 L=1110
recycle 0 plddt 0.563 pae 25.219 rmsd 37.300
recycle 1 plddt 0.571 pae 24.891 rmsd 11.995
recycle 2 plddt 0.567 pae 24.844 rmsd 5.358
recycle 3 plddt 0.567 pae 24.938 rmsd 2.560
recycle 4 plddt 0.567 pae 24.922 rmsd 2.661
recycle 5 plddt 0.561 pae 24.969 rmsd 1.578
recycle 6 plddt 0.562 pae 24.969 rmsd 2.129

**RuntimeError Traceback (most recent call last)
in <cell line: 87>()
90 np.random.seed(seed)
91 npz = f"{jobname}/rf2_seed{seed}_00.npz"
---> 92 pred.predict(inputs=[f"{jobname}/msa.a3m"],
93 out_prefix=f"{jobname}/rf2_seed{seed}",
94 symm=symm,

2 frames
/usr/local/lib/python3.10/dist-packages/torch/functional.py in einsum(*args)
378 # the path for contracting 0 or 1 time(s) is already optimized
379 # or the user has disabled using opt_einsum
--> 380 return _VF.einsum(equation, operands) # type: ignore[attr-defined]
381
382 path = None

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_CUDA_bmm)**

Start with the Colab prediction cell at the pred.predict call shown in the traceback, then inspect the torch.einsum failure for the CPU and cuda:0 operands. Reproduce the default RosettaFold2 dimer run through six recycles using the reported GPU settings. Done means the prediction completes without the device-mismatch RuntimeError.