Using half precision


I am new to using the half-precision for tensors in PyTorch. So I had a very basic question if it’s possible that in my neural network model I can have some variables as half tensors and some as normal full precision tensors?

Basically my model is taking too much memory so instead of decreasing the batch size, I wanted to check if it’s possible to make some variables as half-precision.


Yes, you can change the dtype of a tensor manually and make sure that the performed operations use expected input types.
While this manual approach would be feasible, I would recommend to have a look at the automatic mixed-precision training, which provides utility functions such as automatic casting and gradient scaling.

1 Like


While using amp with my custom forward function using autograd.function. I am getting below error:

NameError: name ‘custom_fwd’ is not defined

I have added @custom_fwd decorator above my forward function.

What does print(torch.cuda.amp.custom_fwd) return?