Sending a tensor (e.g. constant) to device

[muellerzr]

@borgr

Currently, code that worked for me was (for future readers):
l2_reg = torch.tensor(0, dtype=loss.dtype).to(loss.device) # loss some parameter that was in the model hence prepared regularly

Great question! Your recommendation here is actually quite close. prepare doesn’t work with raw tensors, so you do need to change the device on that one manually. You can ensure it will always work by setting it to accelerator.device rather than loss.device