default

Xilinx · Oct 7, 2024 · 0eb2ec2 · 0eb2ec2
1 parent b67fed4
commit 0eb2ec2
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/brevitas/core/scaling/standalone.py b/src/brevitas/core/scaling/standalone.py
@@ -70,10 +70,15 @@ def __init__(
             scaling_init = scaling_init.to(device=device, dtype=dtype)
             if restrict_scaling_impl is not None:
                 scaling_init = restrict_scaling_impl.restrict_init_tensor(scaling_init)
+            else:
+                self.restrict_init_module = Identity()
             self.value = StatelessBuffer(scaling_init.detach())
         else:
             if restrict_scaling_impl is not None:
                 scaling_init = restrict_scaling_impl.restrict_init_float(scaling_init)
+                self.restrict_init_module = restrict_scaling_impl.restrict_init_module()
+            else:
+                self.restrict_init_module = Identity()
             self.value = StatelessBuffer(torch.tensor(scaling_init, dtype=dtype, device=device))
 
     @brevitas.jit.script_method
@@ -149,6 +154,8 @@ def __init__(
         if restrict_scaling_impl is not None:
             scaling_init = restrict_scaling_impl.restrict_init_tensor(scaling_init)
             self.restrict_init_module = restrict_scaling_impl.restrict_init_module()
+        else:
+            self.restrict_init_module = Identity()
         if scaling_init.shape == SCALAR_SHAPE and scaling_shape is not None:
             scaling_init = torch.full(scaling_shape, scaling_init, dtype=dtype, device=device)
         self.value = Parameter(scaling_init)