fix unit test bug in gpu

open-mmlab · Aug 5, 2021 · 74f81e0 · 74f81e0
1 parent 79e7dea
commit 74f81e0
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 12 deletions.
diff --git a/mmedit/models/losses/gan_loss.py b/mmedit/models/losses/gan_loss.py
@@ -111,7 +111,7 @@ def gradient_penalty_loss(discriminator, real_data, fake_data, mask=None):
     """
 
     batch_size = real_data.size(0)
-    alpha = real_data.new_tensor(torch.rand(batch_size, 1, 1, 1))
+    alpha = torch.rand(batch_size, 1, 1, 1).to(real_data)
 
     # interpolate between real_data and fake_data
     interpolates = alpha * real_data + (1. - alpha) * fake_data

diff --git a/tests/test_models/test_inpaintors/test_deepfill_inpaintor.py b/tests/test_models/test_inpaintors/test_deepfill_inpaintor.py
@@ -266,17 +266,6 @@ def test_two_stage_inpaintor():
                 loss_weight=1.0,
             ),
             loss_disc_shift=dict(type='DiscShiftLoss'),
-            loss_composed_percep=dict(
-                type='PerceptualLoss',
-                vgg_type='vgg16',
-                layer_weights={
-                    '4': 1.,
-                    '9': 1.,
-                    '16': 1.,
-                },
-                perceptual_weight=0.05,
-                style_weight=120,
-                pretrained='torchvision://vgg16'),
             pretrained=None)
         tsinpaintor = DeepFillv1Inpaintor(
             **model_, train_cfg=train_cfg, test_cfg=test_cfg).cuda()