[Fix] Fix DCN fp16 (#1014)

* [Fix] Fix DCN fp16 * add comment
2021-05-11 14:16:23 +09:00 · 2021-05-11 14:16:23 +09:00 · 479624672a
parent f61295d944
commit 479624672a
1 changed files with 3 additions and 0 deletions
--- a/mmcv/ops/deform_conv.py
+++ b/mmcv/ops/deform_conv.py
@ -70,6 +70,9 @@ class DeformConv2dFunction(Function):
        ctx.deform_groups = deform_groups
        ctx.im2col_step = im2col_step

+        # until the code is modified for torch.cuda.amp.autocast,
+        # we need to cast weight to avoid type mismatch in fp16 training
+        weight = weight.type_as(input)
        ctx.save_for_backward(input, offset, weight)

        output = input.new_empty(