Can I use pytoch amp functions, GradScaler and autocast on CPU?

Question

I am currently trying to debug my code and would like to run it on the CPU, but I am using torch.cuda.amp.autocast() and torch.cuda.amp.GradScaler(), which are part of the Automatic Mixed Precision package that is from cuda and will be automatically on GPU.

Is there a way to use these functions on the CPU? If not, what alternative approaches can I use to achieve similar functionality on the CPU? Thank you for your help.

Answer 1

You can now have Gradient scaling and autocast on CPU. Use torch.cpu.amp.GradScaler, torch.cpu.amp.autocast for autocast. Your complete code will be something similar to the following

import torch.cpu.amp.GradScaler as scaler
import torch.cpu.amp.autocast as autocast
for batch in train_loader:
    with autocast(device_type="cpu", dtype=torch.bfloat16):
        loss = loss_module(batch, label)
    scaler.scale(loss).backward()
    scaler.unscale_(self.optimizer)
    scaler.step(self.optimizer)
    self.optimizer.zero_grad()