backward compatibility

Author: LeonEthan

nbs/common.modules.ipynb

@@ -509,7 +509,13 @@
     "            A = self.dropout(torch.softmax(scale * scores, dim=-1))\n",
     "            V = torch.einsum(\"bhls,bshd->blhd\", A, values)\n",
     "\n",
-    "            return (V.contiguous(), A) if self.output_attention else (V.contiguous(), None)  "
+    "            return (V.contiguous(), A) if self.output_attention else (V.contiguous(), None)  \n",
+    "\n",
+    "    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):\n",
+    "        # Manually initialize `atten` if not in state_dict\n",
+    "        if prefix + 'atten' not in state_dict:\n",
+    "            self.atten = \"full\"\n",
+    "        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)"
    ]
   },
   {

neuralforecast/common/_modules.py

@@ -429,6 +429,12 @@ def forward(self, queries, keys, values, attn_mask, tau=None, delta=None):
                 (V.contiguous(), A) if self.output_attention else (V.contiguous(), None)
             )
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # Manually initialize `atten` if not in state_dict
+        if prefix + "atten" not in state_dict:
+            self.atten = "full"
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
 # %% ../../nbs/common.modules.ipynb 19
 class PositionalEmbedding(nn.Module):
     def __init__(self, hidden_size, max_len=5000):