Disable autocast in the forward as given in the previous example:
class RoIPool(_RoiPool):
   def forward(self, x, rois):
        with autocast(enabled=False):
        out = super().forward(x.float(), rois) 
    return out

Casts are done internally during the dispatching, but I would recommend not to rely on the internal implementation as it also might depend on the backend (e.g. I don’t know how autocast works exactly on the CPU).
Here is a code snippet showing the input and output dtypes for layers, which can cast to float16 and others which are using float32 due to the needed numerical precision:
def hook(module, input, output):
    print(module)
    print([i.dtype for i in input])
    print([o.dtype for o in output])

lin1 = nn.Linear(10, 10).cuda()    
lin1.register_forward_hook(hook)

lin2 = nn.Linear(10, 10).cuda()    
lin2.register_forward_hook(hook)

sm = nn.Softmax(dim=1).cuda()
sm.register_forward_hook(hook)

x = torch.randn(1, 10).cuda()

with torch.cuda.amp.autocast():
    out = lin1(x)
    out = lin2(out)
    out = sm(out)