Switching device

[enhuiz]

Hi. I notice the model requires passing the device as an argument, which may have not been decided yet at the point of the module initialization. Once the device is decided, it seems we cannot easily change it. Do you consider making the device switchable? One solution may be instead of passing the device, add an attribute:

@property
def device(self):
    return next(self.parameters()).device

[sooftware]

Thank you. Can you send me a pull request?