added docs

This commit is contained in:
Brian Bjarke Jensen
2024-05-21 19:13:24 +02:00
parent 0bbf53d495
commit 6f0513ea2c
+5 -1
View File
@@ -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