diff --git a/onconet/datasets/loader/image.py b/onconet/datasets/loader/image.py
index 32f5dba..e44498b 100644
--- a/onconet/datasets/loader/image.py
+++ b/onconet/datasets/loader/image.py
@@ -161,7 +161,7 @@ def get_image(self, path, additional):
         all_transformers = self.split_transformers[-1][1]
         image = Image.open(path)
         image = apply_transformers_and_cache(image, additional, path, all_transformers,
-                                             self.cache)
+                                             self.cache, cache_full_size=True)
         return image
 
     def get_images(self, paths, additionals):
diff --git a/requirements-dev.txt b/requirements-dev.txt
new file mode 100644
index 0000000..9172e0c
--- /dev/null
+++ b/requirements-dev.txt
@@ -0,0 +1,2 @@
+pytest==7.3.1
+mock==4.0.3
diff --git a/requirements.txt b/requirements.txt
index 1ea5ff7..f6e493f 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,12 +1,9 @@
-torch==1.4.0
-torchvision==0.2.0
-sklearn
-scipy
-numpy
-Pillow
-tqdm
-mock
-scikit-image
-imageio
-gitpython
-lifelines
+imageio==2.28.0
+lifelines==0.27.4
+numpy==1.24.3
+scikit-image==0.20.0
+scikit-learn==1.2.2
+scipy==1.10.1
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.65.0
diff --git a/setup.cfg b/setup.cfg
index 39d6e5d..c4a1f55 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -16,15 +16,12 @@ classifiers =
 packages = find:
 python_requires = >=3
 install_requires =
-    torch==1.4.0
-    torchvision==0.2.0
-    sklearn
-    scipy
-    numpy
-    Pillow
-    tqdm
-    mock
-    scikit-image
-    imageio
-    gitpython
-    lifelines
+    imageio==2.28.0
+    lifelines==0.27.4
+    numpy==1.24.3
+    scikit-image==0.20.0
+    scikit-learn==1.2.2
+    scipy==1.10.1
+    torch==2.0.0
+    torchvision==0.15.1
+    tqdm==4.65.0
diff --git a/tests/test_transformers.py b/tests/test_transformers.py
index f8512bf..b9c7be4 100644
--- a/tests/test_transformers.py
+++ b/tests/test_transformers.py
@@ -75,7 +75,8 @@ def test_image_random_crop(self):
         im.putdata(pixel)
         self.kwargs['w'] = 2
         self.kwargs['h'] = 2
-        with mock.patch('random.randint', lambda x, y: 0):
+        self.args.use_region_annotation = False
+        with mock.patch('torch.randint', lambda x, y, size: torch.tensor([0])):
             cropper = ti.Random_Crop(self.args, self.kwargs)
             output = cropper(im, None)
         expected = Image.new("RGB", (2, 2))
@@ -109,12 +110,12 @@ def test_image_rotate_90(self):
 
     def test_tensor_normalize_tensor_2d(self):
         ''' Test normalizing a tensor with a certain mean and a certain std deviation.'''
-        self.args.img_mean = 2
-        self.args.img_std = 2
-        tensor = torch.IntTensor([[[4, 4, 4], [6, 6, 6]], [[8, 8, 8], [10, 10, 10]]])
+        self.args.img_mean = [2.0]
+        self.args.img_std = [2.0]
+        tensor = torch.FloatTensor([[[4, 4, 4], [6, 6, 6]], [[8, 8, 8], [10, 10, 10]]])
         normalizer = tt.Normalize_Tensor_2d(self.args, self.kwargs)
         output = normalizer(tensor, None).numpy()
-        expected = torch.IntTensor([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]]).numpy()
+        expected = torch.FloatTensor([[[1, 1, 1], [2, 2, 2]], [[3, 3, 3], [4, 4, 4]]]).numpy()
         self.assertTrue(np.array_equal(expected, output))
 
 
diff --git a/tests/test_utils.py b/tests/test_utils.py
index c5031bb..cf683eb 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -13,13 +13,13 @@
 
 import onconet.utils.parsing as parsing
 import onconet.utils.generic as generic
-import onconet.train.state_keeper as state
+import onconet.learn.state_keeper as state
 import datetime
 
 class Test_parse_transformers(unittest.TestCase):
     def setUp(self):
         self.err_msg = "Name of transformer or one of its arguments cant be empty\n\
-                  Use 'name/arg1=value/arg2=value' format"
+                  Use \"name/arg1=value/arg2=value\" format"
 
     def tearDown(self):
         pass
@@ -238,6 +238,7 @@ def setUp(self):
         self.parser.add_argument('--age', default=10)
         self.parser.add_argument('--siblings', default=["alice", "bob"])
         self.parser.add_argument('--save_dir', default=tempfile.mkdtemp())
+        self.parser.add_argument('--use_adv', action='store_true', default=False)
         self.epoch = 10
         self.lr = 0.001
         self.epoch_stats = {}
@@ -256,12 +257,12 @@ def test_simple_model(self):
         model_dict = self.model.state_dict()
         optimizer_dict = self.optimizer.state_dict()
 
-        state_keeper.save(self.model, self.optimizer, self.epoch, self.lr, self.epoch_stats)
+        state_keeper.save({'model': self.model}, {'model': self.optimizer}, self.epoch, self.lr, self.epoch_stats)
         new_model, new_optimizer_state, new_epoch, new_lr, _ = state_keeper.load()
 
         for key in model_dict.keys():
-            self.assertTrue(np.array_equal(model_dict[key].numpy(), new_model.state_dict()[key].numpy()))
-        self.assertEqual(optimizer_dict, new_optimizer_state)
+            self.assertTrue(np.array_equal(model_dict[key].numpy(), new_model['model'].state_dict()[key].numpy()))
+        self.assertEqual(optimizer_dict, new_optimizer_state['model'])
         self.assertEqual(self.epoch, new_epoch)
         self.assertEqual(self.lr, new_lr)