feat: add parameter getters

es-ude · Sep 9, 2023 · 3c184c0 · 3c184c0
1 parent 924e5f3
commit 3c184c0
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/elasticai/creator/nn/fixed_point/linear/layer.py b/elasticai/creator/nn/fixed_point/linear/layer.py
@@ -87,6 +87,22 @@ def __init__(
             device=device,
         )
 
+    @property
+    def lin_weight(self) -> torch.Tensor:
+        return self._linear.weight
+
+    @property
+    def lin_bias(self) -> torch.Tensor | None:
+        return self._linear.bias
+
+    @property
+    def bn_weight(self) -> torch.Tensor:
+        return self._batch_norm.weight
+
+    @property
+    def bn_bias(self) -> torch.Tensor:
+        return self._batch_norm.bias
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         has_batches = x.dim() == 2
         input_shape = x.shape if has_batches else (1, -1)