RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

[SeymourNickelson]

When I try to use this with an exported Torchscript model (from export_torchscript.py) I get the following error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

File "code/torch/piper_train/vits/models.py", line 82, in forward
max_length = ops.prim.NumToTensor(torch.size(x4, 2))
x5 = torch.arange(annotate(number, max_length), dtype=4, layout=None, device=torch.device("cpu"), pin_memory=False)
_23 = torch.lt(torch.unsqueeze(x5, 0), torch.unsqueeze(x_lengths, 1))
~~~~~~~~ <--- HERE

Is there any instructions on how the model should be prepared in order to be used with this repository? Should the entire checkpoint be exported and loaded in? I see the source code for the vits folder is included (with lightning.py) but I don't see pytorch lightning included in the requirements.txt file so I'm wondering how a model should be prepared for use with the sample generator?