Using Automatic Mixed Precision for existing tochvision models

I am trying to use AMP with my resnet50 code from horovod/ at master · horovod/horovod · GitHub. I changed the lines 75 through 86 to the following snippet. As I understand from PyTorch documentation, I am supposed to use autocast in the forward pass. But, since I am using a predefined model (torchvision.models.resnet50()), I don’t have access to that part of the code. So, I am getting the error

File “”, line 78, in train
with autocast:
AttributeError: __enter__

How can I use AMP with these predefined models?

        for i in range(0, len(data), args.batch_size):
            data_batch = data[i:i + args.batch_size]
            target_batch = target[i:i + args.batch_size]
            with autocast:
                output = model(data_batch)
                train_accuracy.update(accuracy(output, target_batch))
                loss = F.cross_entropy(output, target_batch)
                # Average gradients among sub-batches
                loss.div_(math.ceil(float(len(data)) / args.batch_size))
        # Gradient is applied across all ranks

You would have to call the op via:

with autocast():