adding maxpool2d and flatten

DSGT-DLP · Oct 27, 2023 · 6a35554 · 6a35554
1 parent 5cbbaf0
commit 6a35554
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/training/training/core/dl_model.py b/training/training/core/dl_model.py
@@ -17,7 +17,9 @@ class DLModel(nn.Module):
         "SIGMOID": nn.Sigmoid,
         "LOGSOFTMAX": nn.LogSoftmax,
         "CONV2D": nn.Conv2d,
-        "DROPOUT": nn.Dropout
+        "DROPOUT": nn.Dropout,
+        "MAXPOOL2D": nn.MaxPool2d,
+        "FLATTEN": nn.Flatten
     }
 
     def __init__(self, layer_list: list[nn.Module]):