Fix gaitedge, fixes ShiqiYu#91

cocoJennie · Oct 13, 2022 · 41de2d4 · 41de2d4
1 parent c2a4819
commit 41de2d4
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/opengait/modeling/models/gaitedge.py b/opengait/modeling/models/gaitedge.py
@@ -90,11 +90,10 @@ def resume_ckpt(self, restore_hint):
         self._load_ckpt(save_name)
 
     def preprocess(self, sils):
-
-        dilated_mask = (morph.dilation(sils, self.kernel).detach()
+        dilated_mask = (morph.dilation(sils, self.kernel.to(sils.device)).detach()
                         ) > 0.5  # Dilation
-        eroded_mask = (morph.erosion(sils, self.kernel).detach()
-                       ) > 0.5   # Dilation
+        eroded_mask = (morph.erosion(sils, self.kernel.to(sils.device)).detach()
+                       ) > 0.5   # Erosion
         edge_mask = dilated_mask ^ eroded_mask
         return edge_mask, eroded_mask