flatten class added

model.py

@@ -23,6 +23,14 @@ def create_vggface2_model(num_classes:int=2,
     layer_list = list(model_pred.children())[-5:] # all final layers
     model_pred = nn.Sequential(*list(model_pred.children())[:-5])
 
+    class Flatten(nn.Module):
+    def __init__(self):
+        super(Flatten, self).__init__()
+        
+    def forward(self, x):
+        x = x.view(x.size(0), -1)
+        return x
+
     for param in model_pred.parameters():
         param.requires_grad = False