Enhance ScriptModel to support auxiliary head and output extraction from dictionary format

Author: valhassan

geo_deep_learning/tools/script_model.py

@@ -3,17 +3,19 @@
 from torch.nn import functional as F
 
 class ScriptModel(nn.Module):
-    def __init__(self, model, data_module, num_classes, from_logits = True):
+    def __init__(self, model, data_module, num_classes, from_logits = True, aux_head = False):
         super().__init__()
         self.model = model
         self.data_module = data_module
         self.num_classes = num_classes
         self.from_logits = from_logits
-    
+        self.aux_head = aux_head
     def forward(self, x):
         sample = {"image": x}
         preprocessed = self.data_module._model_script_preprocess(sample)
         output = self.model(preprocessed["image"])
+        if isinstance(output, dict):
+            output = output['out']
         if self.from_logits:
             if self.num_classes > 1:
                 output = F.softmax(output, dim=1)