Add abc inheritance/decorators for propper polymorphism.

choderalab · Oct 18, 2023 · 2288017 · 2288017
1 parent 08903e3
commit 2288017
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 4 deletions.
diff --git a/mtenn/combination.py b/mtenn/combination.py
@@ -1,7 +1,9 @@
+import abc
 import torch
 
 
-class Combination(torch.nn.Module):
+class Combination(torch.nn.Module, abc.ABC):
+    @abc.abstractmethod
     def forward(self, pred_list, grad_dict, param_names, *model_params):
         """
         This function signature should be the same for any Combination subclass

diff --git a/mtenn/readout.py b/mtenn/readout.py
@@ -1,8 +1,9 @@
+import abc
 import torch
 from typing import Optional
 
 
-class Readout(torch.nn.Module):
+class Readout(torch.nn.Module, abc.ABC):
     pass
 
 

diff --git a/mtenn/representation.py b/mtenn/representation.py
@@ -1,5 +1,6 @@
+import abc
 import torch
 
 
-class Representation(torch.nn.Module):
+class Representation(torch.nn.Module, abc.ABC):
     pass
diff --git a/mtenn/strategy.py b/mtenn/strategy.py
@@ -1,8 +1,9 @@
+import abc
 from itertools import permutations
 import torch
 
 
-class Strategy(torch.nn.Module):
+class Strategy(torch.nn.Module, abc.ABC):
     pass