One possible approach would be to use the device of another parameter of your model. E.g. assuming your model contains another layer called self.fc:
for i in range(len(sample)):
    data = sample[i].to(self.fc.weight.device)

Alternatively, you could also apply the for loop inside the training loop, but that might need some code changes.