Repository: DAI-Lab/SteganoGAN

Issue:
Refactorize critics and decoders
As we added a new Decoder, we can see that the code it's quiet similar and it should be refactorized in order to follow the same structure as `encoders.py`.
Same it's happening for `critics.py`. Refactorizing this now, will save us time in unit tests and if we want to add more architectures in the future it will be easier.

Discussion:

Patch:
diff --git a/steganogan/critics.py b/steganogan/critics.py
index 0963fd7..301cb61 100644
--- a/steganogan/critics.py
+++ b/steganogan/critics.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 
 class BasicCritic(nn.Module):
@@ -13,29 +13,45 @@ class BasicCritic(nn.Module):
     Output: (N, 1)
     """
 
-    def __init__(self, hidden_size):
-
-        super(BasicCritic, self).__init__()
+    def _conv2d(self, in_channels, out_channels):
+        return nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3
+        )
 
-        self.layers = nn.Sequential(
-            nn.Conv2d(in_channels=3, out_channels=hidden_size, kernel_size=3),
+    def _build_models(self):
+        return nn.Sequential(
+            self._conv2d(3, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=3),
+            self._conv2d(self.hidden_size, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=3),
+            self._conv2d(self.hidden_size, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=1, kernel_size=3)
+            self._conv2d(self.hidden_size, 1)
         )
 
-    def forward(self, x):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.version = '1'
+        self.hidden_size = hidden_size
+        self._models = self._build_models()
+
+    def upgrade_legacy(self):
+        """Transform legacy pretrained models to make them usable with new code versions."""
+        # Transform to version 1
+        if not hasattr(self, 'version'):
+            self._models = self.layers
+            self.version = '1'
 
-        x = self.layers(x)
+    def forward(self, x):
+        x = self._models(x)
         x = torch.mean(x.view(x.size(0), -1), dim=1)
 
         return x
diff --git a/steganogan/decoders.py b/steganogan/decoders.py
index f39828d..f1d150a 100644
--- a/steganogan/decoders.py
+++ b/steganogan/decoders.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 
 class BasicDecoder(nn.Module):
@@ -13,33 +13,62 @@ class BasicDecoder(nn.Module):
     Output: (N, D, H, W)
     """
 
-    def __init__(self, data_depth, hidden_size):
-
-        super(BasicDecoder, self).__init__()
-
-        self.data_depth = data_depth
+    def _conv2d(self, in_channels, out_channels):
+        return nn.Conv2d(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=3,
+            padding=1
+        )
 
+    def _build_models(self):
         self.layers = nn.Sequential(
-            nn.Conv2d(in_channels=3, out_channels=hidden_size, kernel_size=3, padding=1),
+            self._conv2d(3, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=3, padding=1),
+            self._conv2d(self.hidden_size, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=hidden_size, kernel_size=3, padding=1),
+            self._conv2d(self.hidden_size, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size),
 
-            nn.Conv2d(in_channels=hidden_size, out_channels=data_depth, kernel_size=3, padding=1)
+            self._conv2d(self.hidden_size, self.data_depth)
         )
 
+        return [self.layers]
+
+    def __init__(self, data_depth, hidden_size):
+        super().__init__()
+        self.version = '1'
+        self.data_depth = data_depth
+        self.hidden_size = hidden_size
+
+        self._models = self._build_models()
+
+    def upgrade_legacy(self):
+        """Transform legacy pretrained models to make them usable with new code versions."""
+        # Transform to version 1
+        if not hasattr(self, 'version'):
+            self._models = [self.layers]
+
+            self.version = '1'
+
     def forward(self, x):
-        return self.layers(x)
+        x = self._models[0](x)
 
+        if len(self._models) > 1:
+            x_list = [x]
+            for layer in self._models[1:]:
+                x = layer(torch.cat(x_list, dim=1))
+                x_list.append(x)
 
-class DenseDecoder(nn.Module):
+        return x
+
+
+class DenseDecoder(BasicDecoder):
     """
     The DenseDecoder module takes an steganographic image and attempts to decode
     the embedded data tensor.
@@ -47,44 +76,38 @@ class DenseDecoder(nn.Module):
     Input: (N, 3, H, W)
     Output: (N, D, H, W)
     """
-
-    def __init__(self, data_depth, hidden_size):
-        super(DenseDecoder, self).__init__()
-        self.data_depth = data_depth
+    def _build_models(self):
         self.conv1 = nn.Sequential(
-            nn.Conv2d(in_channels=3,
-                      out_channels=hidden_size,
-                      kernel_size=3,
-                      padding=1),
+            self._conv2d(3, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size)
         )
+
         self.conv2 = nn.Sequential(
-            nn.Conv2d(in_channels=hidden_size,
-                      out_channels=hidden_size,
-                      kernel_size=3,
-                      padding=1),
+            self._conv2d(self.hidden_size, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
+            nn.BatchNorm2d(self.hidden_size)
         )
+
         self.conv3 = nn.Sequential(
-            nn.Conv2d(in_channels=hidden_size * 2,
-                      out_channels=hidden_size,
-                      kernel_size=3,
-                      padding=1),
+            self._conv2d(self.hidden_size * 2, self.hidden_size),
             nn.LeakyReLU(inplace=True),
-            nn.BatchNorm2d(hidden_size),
-        )
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(in_channels=hidden_size * 3,
-                      out_channels=data_depth,
-                      kernel_size=3,
-                      padding=1),
+            nn.BatchNorm2d(self.hidden_size)
         )
 
-    def forward(self, x):
-        x1 = self.conv1(x)
-        x2 = self.conv2(x1)
-        x3 = self.conv3(torch.cat((x1, x2), dim=1))
-        x4 = self.conv4(torch.cat((x1, x2, x3), dim=1))
-        return x4
+        self.conv4 = nn.Sequential(self._conv2d(self.hidden_size * 3, self.data_depth))
+
+        return self.conv1, self.conv2, self.conv3, self.conv4
+
+    def upgrade_legacy(self):
+        """Transform legacy pretrained models to make them usable with new code versions."""
+        # Transform to version 1
+        if not hasattr(self, 'version'):
+            self._models = [
+                self.conv1,
+                self.conv2,
+                self.conv3,
+                self.conv4
+            ]
+
+            self.version = '1'
diff --git a/steganogan/encoders.py b/steganogan/encoders.py
index 04cff12..d1a93c8 100644
--- a/steganogan/encoders.py
+++ b/steganogan/encoders.py
@@ -43,10 +43,17 @@ class BasicEncoder(nn.Module):
 
     def __init__(self, data_depth, hidden_size):
         super().__init__()
+        self.version = '1'
         self.data_depth = data_depth
         self.hidden_size = hidden_size
         self._models = self._build_models()
 
+    def upgrade_legacy(self):
+        """Transform legacy pretrained models to make them usable with new code versions."""
+        # Transform to version 1
+        if not hasattr(self, 'version'):
+            self.version = '1'
+
     def forward(self, image, data):
         x = self._models[0](image)
         x_list = [x]
diff --git a/steganogan/models.py b/steganogan/models.py
index b0cc4cf..78afdbc 100644
--- a/steganogan/models.py
+++ b/steganogan/models.py
@@ -345,5 +345,10 @@ class SteganoGAN(object):
         """Loads an instance of SteganoGAN from the given path."""
         steganogan = torch.load(path, map_location='cpu')
         steganogan.verbose = verbose
+
+        steganogan.encoder.upgrade_legacy()
+        steganogan.decoder.upgrade_legacy()
+        steganogan.critic.upgrade_legacy()
+
         steganogan.set_device(cuda)
         return steganogan


Test patch:
diff --git a/tests/test_critics.py b/tests/test_critics.py
new file mode 100644
index 0000000..5779f76
--- /dev/null
+++ b/tests/test_critics.py
@@ -0,0 +1,103 @@
+# -*- coding: utf-8 -*-
+
+import copy
+from unittest import TestCase
+from unittest.mock import Mock, call, patch
+
+import torch
+
+from steganogan import critics
+from tests.utils import assert_called_with_tensors
+
+
+class TestBasicCritic(TestCase):
+
+    class TestCritic(critics.BasicCritic):
+        def __init__(self):
+            pass
+
+    def setUp(self):
+        self.test_critic = self.TestCritic()
+
+    @patch('steganogan.critics.nn.Conv2d', autospec=True)
+    def test__covn2d(self, conv2d_mock):
+        """Conv2d must be called with given args and kernel_size=3 and padding=1"""
+
+        # run
+        result = self.test_critic._conv2d(2, 4)
+
+        # asserts
+        assert result == conv2d_mock.return_value
+        conv2d_mock.assert_called_once_with(
+            in_channels=2,
+            out_channels=4,
+            kernel_size=3,
+        )
+
+    @patch('steganogan.critics.nn.Sequential')
+    @patch('steganogan.critics.nn.BatchNorm2d')
+    @patch('steganogan.critics.nn.Conv2d')
+    def test___init__(self, conv2d_mock, batchnorm2d_mock, sequential_mock):
+        """Test that conv2d and batchnorm are called when creating a new critic with hidden_size"""
+
+        # run
+        critics.BasicCritic(2)
+
+        # assert
+        expected_batch_calls = [call(2), call(2), call(2)]
+        assert batchnorm2d_mock.call_args_list == expected_batch_calls
+
+        expected_conv2d_calls = [
+            call(in_channels=3, out_channels=2, kernel_size=3),
+            call(in_channels=2, out_channels=2, kernel_size=3),
+            call(in_channels=2, out_channels=2, kernel_size=3),
+            call(in_channels=2, out_channels=1, kernel_size=3)
+        ]
+        assert conv2d_mock.call_args_list == expected_conv2d_calls
+
+    def test_forward(self):
+        """Test the return value of method forward"""
+
+        # setup
+        layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+        test_critic = self.TestCritic()
+        test_critic._models = layer1
+
+        # run
+        image = torch.Tensor([[1, 2], [3, 4]])
+        result = test_critic.forward(image)
+
+        # assert
+        expected = torch.Tensor([[5, 6], [7, 8]])
+        expected = torch.mean(expected.view(expected.size(0), -1), dim=1)
+        assert (result == expected).all()
+
+        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
+        assert_called_with_tensors(layer1, [call_1])
+
+    def test_upgrade_legacy_without_version(self):
+        """Test that upgrade legacy works, must set _models to layers"""
+
+        # setup
+        self.test_critic.layers = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+
+        # run
+        self.test_critic.upgrade_legacy()
+
+        # assert
+        assert self.test_critic.version == '1'
+        assert self.test_critic._models == self.test_critic.layers
+
+    @patch('steganogan.critics.nn.Sequential', autospec=True)
+    def test_upgrade_legacy_with_version_1(self, sequential_mock):
+        """The object must be the same and not changed by the method"""
+
+        # setup
+        critic = critics.BasicCritic(1)
+        expected = copy.deepcopy(critic)
+
+        # run
+        critic.upgrade_legacy()
+
+        # assert
+        assert critic.__dict__ == expected.__dict__
diff --git a/tests/test_decoders.py b/tests/test_decoders.py
new file mode 100644
index 0000000..bce9e75
--- /dev/null
+++ b/tests/test_decoders.py
@@ -0,0 +1,190 @@
+# -*- coding: utf-8 -*-
+
+import copy
+from unittest import TestCase
+from unittest.mock import Mock, call, patch
+
+import torch
+
+from steganogan import decoders
+from tests.utils import assert_called_with_tensors
+
+
+class TestBasicDecoder(TestCase):
+
+    class TestDecoder(decoders.BasicDecoder):
+        def __init__(self):
+            pass
+
+    def setUp(self):
+        self.test_decoder = self.TestDecoder()
+
+    @patch('steganogan.decoders.nn.Conv2d', autospec=True)
+    def test__covn2d(self, conv2d_mock):
+        """Conv2d must be called with given args and kernel_size=3 and padding=1"""
+
+        # run
+        result = self.test_decoder._conv2d(2, 4)
+
+        # asserts
+        assert result == conv2d_mock.return_value
+        conv2d_mock.assert_called_once_with(
+            in_channels=2,
+            out_channels=4,
+            kernel_size=3,
+            padding=1
+        )
+
+    @patch('steganogan.decoders.nn.Sequential')
+    @patch('steganogan.decoders.nn.Conv2d')
+    @patch('steganogan.decoders.nn.BatchNorm2d')
+    def test___init__(self, batchnorm_mock, conv2d_mock, sequential_mock):
+        """Test the init params and that the layers are created correctly"""
+
+        # run
+        decoders.BasicDecoder(2, 5)
+
+        # assert
+        expected_batch_calls = [call(5), call(5), call(5)]
+        assert batchnorm_mock.call_args_list == expected_batch_calls
+
+        expected_conv_calls = [
+            call(in_channels=3, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=5, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=5, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=5, out_channels=2, kernel_size=3, padding=1),
+        ]
+        assert conv2d_mock.call_args_list == expected_conv_calls
+
+    def test_upgrade_legacy_without_version(self):
+        """Upgrade legacy must create self._models from conv1, conv2, conv3, conv4"""
+
+        # setup
+        self.test_decoder.layers = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+
+        # run
+        self.test_decoder.upgrade_legacy()
+
+        # assert
+        assert self.test_decoder._models == [self.test_decoder.layers]
+        assert self.test_decoder.version == '1'
+
+    @patch('steganogan.decoders.nn.Sequential', autospec=True)
+    def test_upgrade_legacy_with_version_1(self, sequential_mock):
+        """The object must be the same and not changed by the method"""
+
+        # setup
+        decoder = decoders.BasicDecoder(1, 1)
+        expected = copy.deepcopy(decoder)
+
+        # run
+        decoder.upgrade_legacy()
+
+        # assert
+        assert decoder.__dict__ == expected.__dict__
+
+    def test_forward_1_layer(self):
+        """If there is only one layer it must be called with image as the only argument."""
+
+        # setup
+        layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+        self.test_decoder._models = [layer1]
+
+        # run
+        image = torch.Tensor([[1, 2], [3, 4]])
+        result = self.test_decoder.forward(image)
+
+        # assert
+        assert (result == torch.Tensor([[5, 6], [7, 8]])).all()
+
+        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
+        assert_called_with_tensors(layer1, [call_1])
+
+    def test_forward_more_than_2_layers(self):
+        """If there are more than 2 layers, they must be called adding data to each result"""
+
+        # setup
+        layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+        layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
+        layer3 = Mock(return_value=torch.Tensor([[13, 14], [15, 16]]))
+        self.test_decoder._models = [layer1, layer2, layer3]
+
+        # run
+        image = torch.Tensor([[1, 2], [3, 4]])
+        result = self.test_decoder.forward(image)
+
+        # asserts
+        call_layer_1 = call(torch.Tensor([[1, 2], [3, 4]]))
+        call_layer_2 = call(torch.Tensor([[5, 6], [7, 8]]))
+        call_layer_3 = call(torch.Tensor([[5, 6, 9, 10], [7, 8, 11, 12]]))
+
+        assert_called_with_tensors(layer1, [call_layer_1])
+        assert_called_with_tensors(layer2, [call_layer_2])
+        assert_called_with_tensors(layer3, [call_layer_3])
+
+        assert (result == torch.Tensor([[13, 14], [15, 16]])).all()
+
+
+class TestDenseDecoder(TestCase):
+
+    class TestDecoder(decoders.DenseDecoder):
+        def __init__(self):
+            pass
+
+    def test_upgrade_legacy_without_version(self):
+        """Upgrade legacy must create self._models from conv1, conv2, conv3, conv4"""
+
+        # setup
+        test_decoder = self.TestDecoder()  # instance an empty decoder
+        test_decoder.conv1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
+        test_decoder.conv2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
+        test_decoder.conv3 = Mock(return_value=torch.Tensor([[13, 14], [15, 16]]))
+        test_decoder.conv4 = Mock(return_value=torch.Tensor([[17, 18], [19, 20]]))
+
+        # run
+        test_decoder.upgrade_legacy()
+
+        # assert
+        expected_models = [
+            test_decoder.conv1,
+            test_decoder.conv2,
+            test_decoder.conv3,
+            test_decoder.conv4,
+        ]
+        assert test_decoder._models == expected_models
+        assert test_decoder.version == '1'
+
+    @patch('steganogan.decoders.nn.Sequential', autospec=True)
+    def test_upgrade_legacy_with_version_1(self, sequential_mock):
+        """The object must be the same and not changed by the method"""
+
+        # setup
+        decoder = decoders.DenseDecoder(1, 1)
+        expected = copy.deepcopy(decoder)
+
+        # run
+        decoder.upgrade_legacy()
+
+        # assert
+        assert decoder.__dict__ == expected.__dict__
+
+    @patch('steganogan.decoders.nn.Sequential')
+    @patch('steganogan.decoders.nn.Conv2d')
+    @patch('steganogan.decoders.nn.BatchNorm2d')
+    def test___init__(self, batchnorm_mock, conv2d_mock, sequential_mock):
+        """Test the init params and that the layers are created correctly"""
+
+        # run
+        decoders.DenseDecoder(2, 5)
+
+        # assert
+        expected_batch_calls = [call(5), call(5), call(5)]
+        assert batchnorm_mock.call_args_list == expected_batch_calls
+
+        expected_conv_calls = [
+            call(in_channels=3, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=5, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=10, out_channels=5, kernel_size=3, padding=1),
+            call(in_channels=15, out_channels=2, kernel_size=3, padding=1),
+        ]
+        assert conv2d_mock.call_args_list == expected_conv_calls
diff --git a/tests/test_encoders.py b/tests/test_encoders.py
index ee732c7..3cb685e 100644
--- a/tests/test_encoders.py
+++ b/tests/test_encoders.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 
+import copy
 from unittest import TestCase
 from unittest.mock import Mock, call, patch
 
@@ -34,8 +35,32 @@ class TestBasicEncoder(TestCase):
             padding=1
         )
 
+    def test_upgrade_legacy_no_version(self):
+        """Test that we set a version to our encoder for future code changes"""
+
+        # run
+        self.test_encoder.upgrade_legacy()
+
+        # assert
+        assert self.test_encoder.version == '1'
+
+    @patch('steganogan.encoders.nn.Sequential', autospec=True)
+    def test_upgrade_legacy_with_version_1(self, sequential_mock):
+        """The object must be the same and not changed by the method"""
+
+        # setup
+        encoder = encoders.BasicEncoder(1, 1)
+        expected = copy.deepcopy(encoder)
+
+        # run
+        encoder.upgrade_legacy()
+
+        # assert
+        assert encoder.__dict__ == expected.__dict__
+
     def test_forward_1_layer(self):
         """If there is only one layer it must be called with image as the only argument."""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         self.test_encoder._models = [layer1]
@@ -44,14 +69,15 @@ class TestBasicEncoder(TestCase):
         image = torch.Tensor([[1, 2], [3, 4]])
         result = self.test_encoder.forward(image, 'some_data')
 
-        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
-
         # assert
         assert (result == torch.Tensor([[5, 6], [7, 8]])).all()
+
+        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
         assert_called_with_tensors(layer1, [call_1])
 
     def test_forward_more_than_2_layers(self):
         """If there are more than 2 layers, they must be called adding data to each result"""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
@@ -76,6 +102,7 @@ class TestBasicEncoder(TestCase):
 
     def test_forward_add_image(self):
         """If add_image is true, image must be added to the result."""
+
         # setup
         self.test_encoder.add_image = True
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
@@ -127,6 +154,7 @@ class TestResidualEncoder(TestCase):
 
     def test_forward_1_layer(self):
         """If there is only one layer it must be called with image as the only argument."""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         self.test_encoder._models = [layer1]
@@ -135,14 +163,15 @@ class TestResidualEncoder(TestCase):
         image = torch.Tensor([[1, 2], [3, 4]])
         result = self.test_encoder.forward(image, 'some_data')
 
-        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
-
         # assert
         assert (result == torch.Tensor([[6, 8], [10, 12]])).all()
+
+        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
         assert_called_with_tensors(layer1, [call_1])
 
     def test_forward_more_than_2_layers(self):
         """If there are more than 2 layers, they must be called adding data to each result"""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))
@@ -194,6 +223,7 @@ class TestDenseEncoder(TestCase):
 
     def test_forward_1_layer(self):
         """If there is only one layer it must be called with image as the only argument."""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         self.test_encoder._models = [layer1]
@@ -202,14 +232,15 @@ class TestDenseEncoder(TestCase):
         image = torch.Tensor([[1, 2], [3, 4]])
         result = self.test_encoder.forward(image, 'some_data')
 
-        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
-
         # assert
         assert (result == torch.Tensor([[6, 8], [10, 12]])).all()
+
+        call_1 = call(torch.Tensor([[1, 2], [3, 4]]))
         assert_called_with_tensors(layer1, [call_1])
 
     def test_forward_more_than_2_layers(self):
         """If there are more than 2 layers, they must be called adding data to each result"""
+
         # setup
         layer1 = Mock(return_value=torch.Tensor([[5, 6], [7, 8]]))
         layer2 = Mock(return_value=torch.Tensor([[9, 10], [11, 12]]))