update

pipeline/utils.py

@@ -268,9 +268,9 @@ def DiMeR_reconstruct(model, infer_config, texture_model, texture_model_config,
     # TODO: Device Check
     global_normals = normal_transfer.trans_local_2_global(normals.cpu().permute(0,2,3,1), torch.tensor([0, 90, 180, 270]),
                                                          torch.tensor([5, 5, 5, 5]), radius=4.5,
-                                                         for_lotus=True)
+                                                         for_lotus=True).to(device)
     global_normals = global_normals.permute(0, 3, 1, 2)
-    global_normals = global_normals * multi_view_mask + (1 - multi_view_mask)
+    global_normals = global_normals * multi_view_mask.to(device) + (1 - multi_view_mask.to(device))
     # global_normals = global_normals * fg_mask + bg_mask
     global_normals = F.pad(global_normals, (50, 50, 50, 50), value=1.)
     global_normals = F.interpolate(global_normals, (512, 512), mode='bilinear', align_corners=False)