Compacted everything.

Author: vbharadwaj-bk

openequivariance/openequivariance/_torch/NPDoubleBackwardMixin.py

@@ -36,12 +36,12 @@ def double_backward_cpu(
             grad_outputs=[in1_dgrad_torch, in2_dgrad_torch, weights_dgrad_torch],
         )
 
-        a_np = a.detach().cpu().numpy()
-        b_np = b.detach().cpu().numpy()
-        c_np = c.detach().cpu().numpy()
-        d_np = d.detach().cpu().numpy()
-
-        return (a_np, b_np, c_np, d_np)
+        return (
+            a.detach().cpu().numpy(),
+            b.detach().cpu().numpy(),
+            c.detach().cpu().numpy(),
+            d.detach().cpu().numpy(),
+        )
 
 
 class NumpyDoubleBackwardMixinConv:
@@ -89,9 +89,9 @@ def double_backward_cpu(
             grad_outputs=[in1_dgrad_torch, in2_dgrad_torch, weights_dgrad_torch],
         )
 
-        a_np = a.detach().cpu().numpy()
-        b_np = b.detach().cpu().numpy()
-        c_np = c.detach().cpu().numpy()
-        d_np = d.detach().cpu().numpy()
-
-        return (a_np, b_np, c_np, d_np)
+        return (
+            a.detach().cpu().numpy(),
+            b.detach().cpu().numpy(),
+            c.detach().cpu().numpy(),
+            d.detach().cpu().numpy(),
+        )

openequivariance/openequivariance/_torch/TensorProduct.py

@@ -1,17 +1,17 @@
-import numpy as np
-import torch
-
+from openequivariance.core.LoopUnrollTP import LoopUnrollTP
 from openequivariance import TPProblem
 from openequivariance._torch import extlib
-from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
+import torch
+from openequivariance.core.utils import torch_to_oeq_dtype, dtype_to_enum
+from openequivariance.benchmark.logging_utils import getLogger
 from openequivariance._torch.utils import (
-    enum_to_torch_dtype,
     reorder_torch,
     string_to_tensor,
+    enum_to_torch_dtype,
 )
-from openequivariance.benchmark.logging_utils import getLogger
-from openequivariance.core.LoopUnrollTP import LoopUnrollTP
-from openequivariance.core.utils import dtype_to_enum, torch_to_oeq_dtype
+from openequivariance._torch.NPDoubleBackwardMixin import NumpyDoubleBackwardMixin
+
+import numpy as np
 
 logger = getLogger()