Does it cache both model parameters and model inputs? E.g. if only model parameters were cached, one could wrap the whole inference loop in autocast.
What I also don’t understand is at what moment tensors are cast to fp16? E.g. this is important for implementing forced fp32 module wrappers. E.g. there is an existing external module class that can’t be modified, and we’d like to wrap it / make a new module class that is always fp32 + inputs are not cast to fp16 if they are not coming from prior fp16 processing. E.g. there is a module ROIpool from wetectron/roi_pool.py at master · NVlabs/wetectron · GitHub. It is marked with apex.amp.float_function. I’d like to replace it with torchvision.ops.RoIPool, but how do I wrap ROIPooling module class such that rois argument never gets cast to fp16