Upload model

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e83fde25faac20fb64c1964c7c8f2779059260d11acdc707d71014716880cb77
+oid sha256:65f05a93d49be782d608ddaddd3ed056077922e26890d7acd53b35ad8e7fe540
 size 332799280

modeling_phylogpn.py

@@ -31,12 +31,20 @@ class RCEWeight(nn.Module):
             )
 
         super().__init__()
-        self.input_involution_indices = input_involution_indices
-        self.output_involution_indices = output_involution_indices
+        self._input_involution_indices = input_involution_indices
+        self._output_involution_indices = output_involution_indices
+        self._input_involution_index_tensor = None
+        self._output_involution_index_tensor = None
+        self._device = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        output_involution_indices = torch.tensor(self.output_involution_indices, device=x.device)
-        input_involution_indices = torch.tensor(self.input_involution_indices, device=x.device)
+        if self._device != x.device:
+            self._input_involution_index_tensor = torch.tensor(self._input_involution_indices, device=x.device)
+            self._output_involution_index_tensor = torch.tensor(self._output_involution_indices, device=x.device)
+            self._device = x.device
+
+        output_involution_indices = self._output_involution_index_tensor
+        input_involution_indices = self._input_involution_index_tensor
         return (x + x[output_involution_indices][:, input_involution_indices].flip(2)) / 2
 
 
@@ -46,10 +54,16 @@ class IEBias(nn.Module):
             raise ValueError("`involution_indices` must be an involution")
 
         super().__init__()
-        self.involution_indices = involution_indices
+        self._involution_indices = involution_indices
+        self._involution_index_tensor = None
+        self._device = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        involution_indices = torch.tensor(self.involution_indices, device=x.device)
+        if self._device != x.device:
+            self._involution_index_tensor = torch.tensor(self._involution_indices, device=x.device)
+            self._device = x.device
+
+        involution_indices = self._involution_index_tensor
         return (x + x[involution_indices]) / 2
 
 
@@ -64,23 +78,25 @@ class IEWeight(nn.Module):
             )
 
         super().__init__()
-        self.input_involution_indices = input_involution_indices
-        self.output_involution_indices = output_involution_indices
+        self._input_involution_indices = input_involution_indices
+        self._output_involution_indices = output_involution_indices
+        self._input_involution_index_tensor = None
+        self._output_involution_index_tensor = None
+        self._device = None
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        input_involution_indices = torch.tensor(self.input_involution_indices, device=x.device)
-        output_involution_indices = torch.tensor(self.output_involution_indices, device=x.device)
+        if self._device != x.device:
+            self._input_involution_index_tensor = torch.tensor(self._input_involution_indices, device=x.device)
+            self._output_involution_index_tensor = torch.tensor(self._output_involution_indices, device=x.device)
+            self._device = x.device
+
+        output_involution_indices = self._output_involution_index_tensor
+        input_involution_indices = self._input_involution_index_tensor
         return (x + x[input_involution_indices][:, output_involution_indices]) / 2
-    
+
 
 class RCEByteNetBlock(nn.Module):
-    def __init__(
-        self,
-        outer_involution_indices: List[int],
-        inner_dim: int,
-        kernel_size: int,
-        dilation_rate: int = 1
-    ):
+    def __init__(self, outer_involution_indices: List[int], inner_dim: int, kernel_size: int, dilation_rate: int = 1):
         outer_dim = len(outer_involution_indices)
 
         if outer_dim % 2 != 0:
@@ -130,7 +146,6 @@ class RCEByteNetBlock(nn.Module):
             layers[8], "bias",
             IEBias(outer_involution_indices)
         )
-
         self.layers = nn.Sequential(*layers)
         self._kernel_size = kernel_size
         self._dilation_rate = dilation_rate