Tensors of the same index must be on the same device and the same dtype except `step` tensors that can be CPU and float32 notwithstanding

torch.set_default_dtype(torch.float32)