If you are creating tensors inside your forward pass, Please use the below code
class Model(nn.Module):
    def __init__(self, args):
        <insert initialization code>
    def forward(self, original, keypoints):
        some_tensor = torch.tensor([....], device=original.device)
        <do some thing with some_tensor>