[FIX] Check LSTM proj_size only for torch>=1.8.0

Resolves fKunstner/backpack-discuss#112.
f-dangel · Jun 30, 2021 · 4f126b0 · 4f126b0
1 parent 5726cc1
commit 4f126b0
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/backpack/core/derivatives/lstm.py b/backpack/core/derivatives/lstm.py
@@ -5,6 +5,7 @@
 from torch.nn import LSTM
 
 from backpack.core.derivatives.basederivatives import BaseParameterDerivatives
+from backpack.utils import TORCH_VERSION, VERSION_1_8_0
 
 
 class LSTMDerivatives(BaseParameterDerivatives):
@@ -47,8 +48,9 @@ def _check_parameters(module: LSTM) -> None:
             raise NotImplementedError("only dropout = 0 is supported")
         if module.bidirectional is not False:
             raise NotImplementedError("only bidirectional = False is supported")
-        if module.proj_size != 0:
-            raise NotImplementedError("only proj_size = 0 is supported")
+        if TORCH_VERSION >= VERSION_1_8_0:
+            if module.proj_size != 0:
+                raise NotImplementedError("only proj_size = 0 is supported")
 
     @staticmethod
     def _forward_pass(

diff --git a/backpack/utils/__init__.py b/backpack/utils/__init__.py
@@ -1 +1,6 @@
 """Contains utility functions."""
+
+from pkg_resources import get_distribution, packaging
+
+TORCH_VERSION = packaging.version.parse(get_distribution("torch").version)
+VERSION_1_8_0 = packaging.version.parse("1.8.0")