Questions about torch.cuda.amp

1

To train my model using less GPU memory, I tried to use the “amp” package. However,an error occurred when importing this package, which is:

"ModuleNotFoundError: No module named 'torch.cuda.amp.autocast'".

My pytorch verision is 1.8.1, and the codes for importting the package are written as:

import torch.cuda.amp.autocast
import torch.cuda.amp.GradScaler

I wonder why this error occurred, and how to solve it. Thanks

2

You could directly use these methods and objects via:

import torch
torch.cuda.amp.autocast
torch.cuda.amp.GradScaler

or could import them via:

from torch.cuda.amp import autocast
from torch.cuda.amp import GradScaler
1 Like