Could you print the device of self.positions, input, and self.mask before the line of code which throws the error?
Skimming through your code I cannot find any obvious mistake where you force a certain device.