torch 2 inf

michalwols · Aug 30, 2023 · 27a5f35 · 27a5f35
1 parent 7a02e18
commit 27a5f35
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/yann/__init__.py b/yann/__init__.py
@@ -446,7 +446,11 @@ def grad_norm(parameters, norm_type: float = 2.0):
   if len(parameters) == 0:
     return torch.tensor(0.)
   device = parameters[0].grad.device
-  from torch._six import inf
+  try:
+    from torch import inf
+  except ImportError:
+    from torch._six import inf
+
   if norm_type == inf:
     norms = [p.grad.detach().abs().max().to(device) for p in parameters]
     norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
@@ -471,7 +475,10 @@ def param_norm(parameters, norm_type: float = 2.0):
   if len(parameters) == 0:
     return torch.tensor(0.)
   device = parameters[0].device
-  from torch._six import inf
+  try:
+    from torch import inf
+  except ImportError:
+    from torch._six import inf
   if norm_type == inf:
     norms = [p.detach().abs().max().to(device) for p in parameters]
     norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))

diff --git a/yann/callbacks/callbacks.py b/yann/callbacks/callbacks.py
@@ -104,7 +104,7 @@ def decorated(func):
       return decorated
 
   @callback
-  def on_init(self, trainer=None):
+  def on_init(self, trainer=None, kwargs=None):
     pass
 
   @callback