MaskRCNN and Apex O2 issue

I am trying to use O2 but roialign is giving me the following error
"File “/opt/conda/lib/python3.7/site-packages/torchvision/ops/”, line 45, in roi_align
sampling_ratio, aligned)
RuntimeError: Expected tensor for argument #1 ‘input’ to have the same type as tensor for argument #2 ‘rois’; but type torch.cuda.HalfTensor does not equal torch.cuda.FloatTensor (while checking arguments for ROIAlign_forward_cuda)
Can you please guide me how to resolve this error?


Double post from here.