added device property for autoencoder (lucidrains#86)

MarcusLoppe · Jun 17, 2024 · 9060ef1 · 9060ef1
1 parent 019b5f5
commit 9060ef1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/meshgpt_pytorch/meshgpt_pytorch.py b/meshgpt_pytorch/meshgpt_pytorch.py
@@ -635,6 +635,10 @@ def __init__(
         self.commit_loss_weight = commit_loss_weight
         self.bin_smooth_blur_sigma = bin_smooth_blur_sigma
 
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
     @classmethod
     def _from_pretrained(
         cls,