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/roi_align.py”, 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?

Thanks
Walid

Double post from here.