optimised

architectures/WideResNet.py

@@ -0,0 +1,15 @@
+
+
+import torch.nn as nn
+from torchvision import models
+
+# Base model
+from architectures.Model import Model
+
+class WideResNet(Model):
+    
+    def get(self):
+        # wide_resnet50_2 is a common high-performance choice
+        m = models.wide_resnet50_2(weights=models.Wide_ResNet50_2_Weights.DEFAULT)
+        m.fc = nn.Linear(m.fc.in_features, self.size)
+        return m