model/src/models/visual_communication.py

@@ -4,8 +4,6 @@ from __future__ import annotations
 
 from torch import nn
 
-from shared.docstore.src.classes import ModelData
-
 from .angle import AngleTail
 from .contact import ContactTail
 from .distance import DistanceTail
@@ -39,7 +37,7 @@ class VisualCommunicationModel(nn.Module):
         self.framing_tail = FramingTail()
         self.salience_tail = SalienceTail()
 
-    def forward(self, x) -> ModelData:
+    def forward(self, x) -> dict:
         """Calculate model output on data."""
         # generate visual representation
         features = self.resnet_head(x)
@@ -57,6 +55,4 @@ class VisualCommunicationModel(nn.Module):
             'framing': self.framing_tail(features),
             'salience': self.salience_tail(features),
         }
-        # convert to respective classes
-        data = ModelData.from_prediction_dict(prediction_dict)
-        return data
+        return prediction_dict