Update model.py

model.py

@@ -338,5 +338,5 @@ class MultiHeadClassification(nn.Module):
         backbone = AutoModel.from_pretrained(os.path.join(model_path, 'pretrained/backbone'))
         instance = cls(backbone, head_config, dropout, l2_reg)
         instance.load(os.path.join(model_path, 'multi-head-sequence-classification-model-model.pth'))
-        instance.head_config = {k: v.weights.shape[1] for k, v in instance.heads.items()}
+        instance.head_config = {k: v.weight.shape[1] for k, v in instance.heads.items()}
         return instance