perf: clear buffers

junhsss · Feb 23, 2023 · 2feb1ad · 2feb1ad
1 parent c83c07e
commit 2feb1ad
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/einorm/__init__.py b/einorm/__init__.py
@@ -2,4 +2,4 @@
 
 __all__ = ["Einorm"]
 
-__version__ = "0.1.0"
+__version__ = "0.1.1"
diff --git a/einorm/einorm.py b/einorm/einorm.py
@@ -98,18 +98,18 @@ def __init__(
         _perm = torch.tensor(
             _perm_dims,
             dtype=torch.long,
-            device=device,
         )
 
+        self.perm = _perm
+
         _inv = torch.empty_like(_perm)
         _inv[_perm] = torch.arange(_perm.size(0))
 
+        self.inv = _inv
+
         # skip permute whenever possible
         self.skip_perm = _perm_dims == list(range(len(_pattern)))
 
-        self.register_buffer("perm", _perm)
-        self.register_buffer("inv", _inv)
-
         group_shape: Tuple[int, ...] = ()
         if group:
             for axis in _group:

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "einorm"
-version = "0.1.0"
+version = "0.1.1"
 description = "An einops-style generalized normalization layer."
 authors = ["junhsss <junhsssr@gmail.com>"]
 readme = "README.md"
Original file line number	Diff line number	Diff line change
Expand Up		@@ -2,4 +2,4 @@

		__all__ = ["Einorm"]

		__version__ = "0.1.0"
		__version__ = "0.1.1"