display weighted module dtype and device (#173)

Co-authored-by: Benjamin Trom <[email protected]>
finegrain-ai · Jan 11, 2024 · 457c3f5 · 457c3f5
1 parent 14ce2f5
commit 457c3f5
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/refiners/fluxion/layers/module.py b/src/refiners/fluxion/layers/module.py
@@ -164,6 +164,9 @@ def device(self) -> Device:
     def dtype(self) -> DType:
         return self.weight.dtype
 
+    def __str__(self) -> str:
+        return f"{super().__str__().removesuffix(')')}, device={self.device}, dtype={str(self.dtype).removeprefix('torch.')})"
+
 
 class TreeNode(TypedDict):
     value: str