fix torch version check for resolve_conj

state-spaces · Mar 18, 2022 · 83a9f13 · 83a9f13 · DorotheaKolossa · Mar 18, 2022
1 parent 07a0c61
commit 83a9f13
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/models/sequence/ss/kernel.py b/src/models/sequence/ss/kernel.py
@@ -62,7 +62,7 @@
 _c2r = torch.view_as_real
 _r2c = torch.view_as_complex
 
-if torch.__version__.startswith('1.10'):
+if tuple(map(int, torch.__version__.split('.')[:2])) >= (1, 10):
     _resolve_conj = lambda x: x.conj().resolve_conj()
 else:
     _resolve_conj = lambda x: x.conj()

diff --git a/src/models/sequence/ss/standalone/s4.py b/src/models/sequence/ss/standalone/s4.py
@@ -94,7 +94,7 @@ def _broadcast_dims(*tensors):
 _c2r = torch.view_as_real
 _r2c = torch.view_as_complex
 _conj = lambda x: torch.cat([x, x.conj()], dim=-1)
-if torch.__version__.startswith('1.10'):
+if tuple(map(int, torch.__version__.split('.')[:2])) >= (1, 10):
     _resolve_conj = lambda x: x.conj().resolve_conj()
 else:
     _resolve_conj = lambda x: x.conj()