I’m importing a module class, so I’d like to make a wrapping module class so that autocast wouldn’t cast unprocessed inputs to this module as fp16. Can I achieve that?
Pseudocode:
from torchvision.ops.misc import RoIPool as _RoIPool

class RoIPool(_RoiPool):
   def forward(self, x, rois)
     # how to decorate this function so that rois is never converted to fp16?
     return super().forward(x.float(), rois).type_as(x)

# MyModel processes x and does not process rois (pipes them directly to the module)
model = MyModel(RoiPool())

for x, rois in data_loader:
  with autocast(enabled = True):
    model(x, rois).backward()
  ...

Part of my problem is that I don’t understand how autocast works: at what point / how are casts done?