From 6ff41d58b9e8f8b2e4af47f8f60ffb7b628a49c5 Mon Sep 17 00:00:00 2001
From: Varuna Jayasiri <vpjayasiri@gmail.com>
Date: Thu, 19 Aug 2021 15:48:49 +0530
Subject: [PATCH] __call__ -> forward

---
 labml_nn/graphs/gat/__init__.py      |  2 +-
 labml_nn/graphs/gat/experiment.py    |  2 +-
 labml_nn/graphs/gatv2/__init__.py    |  2 +-
 labml_nn/graphs/gatv2/experiment.py  |  2 +-
 labml_nn/hypernetworks/experiment.py |  2 +-
 labml_nn/hypernetworks/hyper_lstm.py | 11 ++++++-----
 6 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/labml_nn/graphs/gat/__init__.py b/labml_nn/graphs/gat/__init__.py
index 218faa4c..76b91df5 100644
--- a/labml_nn/graphs/gat/__init__.py
+++ b/labml_nn/graphs/gat/__init__.py
@@ -85,7 +85,7 @@ class GraphAttentionLayer(Module):
         # Dropout layer to be applied for attention
         self.dropout = nn.Dropout(dropout)
 
-    def __call__(self, h: torch.Tensor, adj_mat: torch.Tensor):
+    def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
         """
         * `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
         * `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
diff --git a/labml_nn/graphs/gat/experiment.py b/labml_nn/graphs/gat/experiment.py
index 903c9acf..b2c9a41f 100644
--- a/labml_nn/graphs/gat/experiment.py
+++ b/labml_nn/graphs/gat/experiment.py
@@ -134,7 +134,7 @@ class GAT(Module):
         # Dropout
         self.dropout = nn.Dropout(dropout)
 
-    def __call__(self, x: torch.Tensor, adj_mat: torch.Tensor):
+    def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
         """
         * `x` is the features vectors of shape `[n_nodes, in_features]`
         * `adj_mat` is the adjacency matrix of the form
diff --git a/labml_nn/graphs/gatv2/__init__.py b/labml_nn/graphs/gatv2/__init__.py
index ae1befe7..8e6e2291 100644
--- a/labml_nn/graphs/gatv2/__init__.py
+++ b/labml_nn/graphs/gatv2/__init__.py
@@ -121,7 +121,7 @@ class GraphAttentionV2Layer(Module):
         # Dropout layer to be applied for attention
         self.dropout = nn.Dropout(dropout)
 
-    def __call__(self, h: torch.Tensor, adj_mat: torch.Tensor):
+    def forward(self, h: torch.Tensor, adj_mat: torch.Tensor):
         """
         * `h`, $\mathbf{h}$ is the input node embeddings of shape `[n_nodes, in_features]`.
         * `adj_mat` is the adjacency matrix of shape `[n_nodes, n_nodes, n_heads]`.
diff --git a/labml_nn/graphs/gatv2/experiment.py b/labml_nn/graphs/gatv2/experiment.py
index 76d7e58a..30889c6d 100644
--- a/labml_nn/graphs/gatv2/experiment.py
+++ b/labml_nn/graphs/gatv2/experiment.py
@@ -50,7 +50,7 @@ class GATv2(Module):
         # Dropout
         self.dropout = nn.Dropout(dropout)
 
-    def __call__(self, x: torch.Tensor, adj_mat: torch.Tensor):
+    def forward(self, x: torch.Tensor, adj_mat: torch.Tensor):
         """
         * `x` is the features vectors of shape `[n_nodes, in_features]`
         * `adj_mat` is the adjacency matrix of the form
diff --git a/labml_nn/hypernetworks/experiment.py b/labml_nn/hypernetworks/experiment.py
index 0b3e7c1b..64989a7c 100644
--- a/labml_nn/hypernetworks/experiment.py
+++ b/labml_nn/hypernetworks/experiment.py
@@ -22,7 +22,7 @@ class AutoregressiveModel(Module):
         self.lstm = rnn_model
         self.generator = nn.Linear(d_model, n_vocab)
 
-    def __call__(self, x: torch.Tensor):
+    def forward(self, x: torch.Tensor):
         x = self.src_embed(x)
         # Embed the tokens (`src`) and run it through the the transformer
         res, state = self.lstm(x)
diff --git a/labml_nn/hypernetworks/hyper_lstm.py b/labml_nn/hypernetworks/hyper_lstm.py
index ae911858..b565130e 100644
--- a/labml_nn/hypernetworks/hyper_lstm.py
+++ b/labml_nn/hypernetworks/hyper_lstm.py
@@ -147,9 +147,9 @@ class HyperLSTMCell(Module):
         self.layer_norm = nn.ModuleList([nn.LayerNorm(hidden_size) for _ in range(4)])
         self.layer_norm_c = nn.LayerNorm(hidden_size)
 
-    def __call__(self, x: torch.Tensor,
-                 h: torch.Tensor, c: torch.Tensor,
-                 h_hat: torch.Tensor, c_hat: torch.Tensor):
+    def forward(self, x: torch.Tensor,
+                h: torch.Tensor, c: torch.Tensor,
+                h_hat: torch.Tensor, c_hat: torch.Tensor):
         # $$
         # \hat{x}_t = \begin{pmatrix}
         # h_{t-1} \\
@@ -202,6 +202,7 @@ class HyperLSTM(Module):
     """
     # HyperLSTM module
     """
+
     def __init__(self, input_size: int, hidden_size: int, hyper_size: int, n_z: int, n_layers: int):
         """
         Create a network of `n_layers` of HyperLSTM.
@@ -220,8 +221,8 @@ class HyperLSTM(Module):
                                    [HyperLSTMCell(hidden_size, hidden_size, hyper_size, n_z) for _ in
                                     range(n_layers - 1)])
 
-    def __call__(self, x: torch.Tensor,
-                 state: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = None):
+    def forward(self, x: torch.Tensor,
+                state: Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]] = None):
         """
         * `x` has shape `[n_steps, batch_size, input_size]` and
         * `state` is a tuple of $h, c, \hat{h}, \hat{c}$.