Align eps dimension to previous fix

pulp-platform · Feb 23, 2022 · 5ea3338 · 5ea3338
1 parent 57d5f99
commit 5ea3338
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/transf/deploy.py b/nemo/transf/deploy.py
@@ -106,9 +106,9 @@ def _set_eps_in_pact(self, eps_in):
             if eps_in_new is None:
                 continue
         if (m.__class__.__name__ == "PACT_Act"):
-            m.eps_in = eps_in_new[0].clone().detach().requires_grad_(False)
+            m.eps_in = eps_in_new.clone().detach().requires_grad_(False)
         if (m.__class__.__name__ == "PACT_QuantizedBatchNormNd"):
-            m.eps_in = eps_in_new[0].clone().detach().requires_grad_(False)
+            m.eps_in = eps_in_new.clone().detach().requires_grad_(False)
         if (m.__class__.__name__ == "PACT_IntegerAdd"):
             eps_in_list = []
             for eps in eps_in_new: