model/src/models/visual_communication.py

@@ -1,6 +1,7 @@
+"""Definition of VisualCommunicationModel class."""
 from __future__ import annotations
 
-import torch.nn as nn
+from torch import nn
 
 from .angle import AngleTail
 from .contact import ContactTail
@@ -28,6 +29,8 @@ from shared.dto import VisualSyntaxData
 
 
 class VisualCommunicationModel(nn.Module):
+    """Visual communication model."""
+
     def __init__(self):
         super().__init__()
         # store other models
@@ -45,6 +48,7 @@ class VisualCommunicationModel(nn.Module):
         self.salience_tail = SalienceTail()
 
     def forward(self, x):
+        """Calculate model output on data."""
         # generate visual representation
         vis_rep = self.resnet_head(x)
         # prepare result map