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

Probably a beginner’s question:

I couldn’t find a simple way to send a tensor to the device. If I have a constant (e.g. 0, or frozen weights for l2 regularization) Do I have to add them to a model for the accelerator to send them to a device (through prepare)

I would expect for example:
const = accelerator.prepare(torch.tensor(0.))

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 :slight_smile:

1 Like