diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 1068b09..4c03425 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,8 +1,9 @@
 [bumpversion]
-current_version = 2.1.0
+current_version = 2.1.1
 commit = False
 tag = True
 
 [bumpversion:file:setup.py]
 
 [bumpversion:file:ISR/__init__.py]
+
diff --git a/ISR/__init__.py b/ISR/__init__.py
index 44f7224..f947288 100644
--- a/ISR/__init__.py
+++ b/ISR/__init__.py
@@ -1,3 +1,3 @@
 from . import assistant
 
-__version__ = '2.1.0'
+__version__ = '2.1.1'
diff --git a/ISR/models/cut_vgg19.py b/ISR/models/cut_vgg19.py
index 05fc030..707f81b 100644
--- a/ISR/models/cut_vgg19.py
+++ b/ISR/models/cut_vgg19.py
@@ -1,6 +1,5 @@
-from keras.models import Model
-from keras.applications.vgg19 import VGG19
-from keras.layers import Input
+from tensorflow.keras.models import Model
+from tensorflow.keras.applications.vgg19 import VGG19
 from ISR.utils.logger import get_logger
 
 
@@ -36,15 +35,9 @@ def _cut_vgg(self):
         layers selected by self.layers_to_extract.
         """
 
-        vgg = VGG19(weights='imagenet', include_top=False)
+        vgg = VGG19(weights='imagenet', include_top=False, input_shape=self.input_shape)
         vgg.trainable = False
-        collect = []
-        for i in self.layers_to_extract:
-            collect.append(vgg.layers[i].output)
-
-        vgg.outputs = collect
-        hr = Input(shape=self.input_shape)
-        features = vgg(hr)
-        self.model = Model(inputs=hr, outputs=features)
-        self.model.name = 'feature_extractor'
+        outputs = [vgg.layers[i].output for i in self.layers_to_extract]
+        self.model = Model([vgg.input], outputs)
+        self.model._name = 'feature_extractor'
         self.name = 'vgg19'  # used in weights naming
diff --git a/ISR/models/discriminator.py b/ISR/models/discriminator.py
index 797a681..44678ff 100644
--- a/ISR/models/discriminator.py
+++ b/ISR/models/discriminator.py
@@ -1,7 +1,7 @@
-from keras.layers import concatenate, Flatten, Input, Activation, Dense, Conv2D, BatchNormalization
-from keras.layers.advanced_activations import LeakyReLU
-from keras.models import Model
-from keras.optimizers import Adam
+import tensorflow as tf
+from tensorflow.keras.layers import concatenate, Flatten, Input, Activation, Dense, Conv2D, BatchNormalization, LeakyReLU
+from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
 
 
 class Discriminator:
@@ -16,7 +16,7 @@ class Discriminator:
     Attributes:
         model: Keras model.
         name: name used to identify what discriminator is used during GANs training.
-        model.name: identifies this network as the discriminator network
+        model._name: identifies this network as the discriminator network
             in the compound model built by the trainer class.
         block_param: dictionary, determines the number of filters and the strides for each
             conv block.
@@ -33,7 +33,7 @@ def __init__(self, patch_size, kernel_size=3):
         self.model = self._build_disciminator()
         optimizer = Adam(0.0002, 0.5)
         self.model.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy'])
-        self.model.name = 'discriminator'
+        self.model._name = 'discriminator'
         self.name = 'srgan-large'
 
     def _conv_block(self, input, filters, strides, batch_norm=True, count=None):
diff --git a/ISR/models/rdn.py b/ISR/models/rdn.py
index 2db086b..698c0e1 100644
--- a/ISR/models/rdn.py
+++ b/ISR/models/rdn.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
-from keras.initializers import RandomUniform
-from keras.layers import concatenate, Input, Activation, Add, Conv2D, Lambda, UpSampling2D
-from keras.models import Model
+from tensorflow.keras.initializers import RandomUniform
+from tensorflow.keras.layers import concatenate, Input, Activation, Add, Conv2D, Lambda, UpSampling2D
+from tensorflow.keras.models import Model
 from ISR.models.imagemodel import ImageModel
 
 
@@ -37,7 +37,7 @@ class RDN(ImageModel):
         x: integer, the scaling factor.
         model: Keras model of the RDN.
         name: name used to identify what upscaling network is used during training.
-        model.name: identifies this network as the generator network
+        model._name: identifies this network as the generator network
             in the compound model built by the trainer class.
     """
 
@@ -64,7 +64,7 @@ def __init__(
             minval=-init_extreme_val, maxval=init_extreme_val, seed=None
         )
         self.model = self._build_rdn()
-        self.model.name = 'generator'
+        self.model._name = 'generator'
         self.name = 'rdn'
 
     def _upsampling_block(self, input_layer):
@@ -90,7 +90,7 @@ def _pixel_shuffle(self, input_layer):
             kernel_initializer=self.initializer,
         )(input_layer)
         return Lambda(
-            lambda x: tf.depth_to_space(x, block_size=self.scale, data_format='NHWC'),
+            lambda x: tf.nn.depth_to_space(x, block_size=self.scale, data_format='NHWC'),
             name='PixelShuffle',
         )(x)
 
diff --git a/ISR/models/rrdn.py b/ISR/models/rrdn.py
index 257161d..53d6df6 100644
--- a/ISR/models/rrdn.py
+++ b/ISR/models/rrdn.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
-from keras.initializers import RandomUniform
-from keras.layers import UpSampling2D, concatenate, Input, Activation, Add, Conv2D, Lambda
-from keras.models import Model
+from tensorflow.keras.initializers import RandomUniform
+from tensorflow.keras.layers import UpSampling2D, concatenate, Input, Activation, Add, Conv2D, Lambda
+from tensorflow.keras.models import Model
 from ISR.models.imagemodel import ImageModel
 
 
@@ -39,7 +39,7 @@ class RRDN(ImageModel):
         x: integer, the scaling factor.
         model: Keras model of the RRDN.
         name: name used to identify what upscaling network is used during training.
-        model.name: identifies this network as the generator network
+        model._name: identifies this network as the generator network
             in the compound model built by the trainer class.
     """
 
@@ -59,7 +59,7 @@ def __init__(
         self.kernel_size = kernel_size
         self.patch_size = patch_size
         self.model = self._build_rdn()
-        self.model.name = 'generator'
+        self.model._name = 'generator'
         self.name = 'rrdn'
 
     def _dense_block(self, input_layer, d, t):
@@ -122,7 +122,7 @@ def _pixel_shuffle(self, input_layer):
             name='PreShuffle',
         )(input_layer)
         return Lambda(
-            lambda x: tf.depth_to_space(x, block_size=self.scale, data_format='NHWC'),
+            lambda x: tf.nn.depth_to_space(x, block_size=self.scale, data_format='NHWC'),
             name='PixelShuffle',
         )(x)
 
diff --git a/ISR/train/trainer.py b/ISR/train/trainer.py
index 2ca830a..1cf650b 100644
--- a/ISR/train/trainer.py
+++ b/ISR/train/trainer.py
@@ -1,12 +1,11 @@
-import yaml
 import numpy as np
 from time import time
 from tqdm import tqdm
-from keras.models import Model
-from keras.layers import Input
-from keras.optimizers import Adam
-from keras.callbacks import TensorBoard
-from keras import backend as K
+from tensorflow.keras.models import Model
+from tensorflow.keras.layers import Input
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.callbacks import TensorBoard
+from tensorflow.keras import backend as K
 from ISR.utils.datahandler import DataHandler
 from ISR.utils.train_helper import TrainerHelper
 from ISR.utils.metrics import PSNR
@@ -292,7 +291,7 @@ def train(self, epochs, steps_per_epoch, batch_size, monitored_metrics):
             self
         )  # load_weights, creates folders, creates basename
 
-        self.tensorboard = TensorBoard(log_dir=self.helper.callback_paths['logs'])
+        self.tensorboard = TensorBoard(log_dir=str(self.helper.callback_paths['logs']))
         self.tensorboard.set_model(self.model)
 
         # validation data
diff --git a/ISR/utils/metrics.py b/ISR/utils/metrics.py
index 36da298..79a42cf 100644
--- a/ISR/utils/metrics.py
+++ b/ISR/utils/metrics.py
@@ -1,4 +1,4 @@
-import keras.backend as K
+import tensorflow.keras.backend as K
 
 
 def PSNR(y_true, y_pred, MAXp=1):
diff --git a/ISR/utils/train_helper.py b/ISR/utils/train_helper.py
index 715c12e..11d0a69 100644
--- a/ISR/utils/train_helper.py
+++ b/ISR/utils/train_helper.py
@@ -181,7 +181,7 @@ def _save_weights(self, epoch, generator, discriminator=None, metric=None, best=
                 (self.weights_name['generator'].name).format(metric='', epoch=epoch + 1)
             )
         # CANT SAVE MODEL DUE TO TF LAYER INSIDE LAMBDA (PIXELSHUFFLE)
-        generator.save_weights(gen_path)
+        generator.save_weights(gen_path.as_posix())
         if discriminator:
             if best:
                 discr_path = self.weights_name['discriminator'].with_name(
@@ -193,7 +193,7 @@ def _save_weights(self, epoch, generator, discriminator=None, metric=None, best=
                 discr_path = self.weights_name['discriminator'].with_name(
                     (self.weights_name['discriminator'].name).format(metric='', epoch=epoch + 1)
                 )
-            discriminator.model.save_weights(discr_path)
+            discriminator.model.save_weights(discr_path.as_posix())
         try:
             self._remove_old_weights(self.max_n_other_weights, max_best=self.max_n_best_weights)
         except Exception as e:
diff --git a/setup.py b/setup.py
index 6631e03..39f35a3 100644
--- a/setup.py
+++ b/setup.py
@@ -10,17 +10,17 @@
 
 setup(
     name='ISR',
-    version='2.1.0',
+    version='2.1.1',
     author='Francesco Cardinale',
     author_email='testadicardi@gmail.com',
     description='Image Super Resolution',
     long_description=long_description,
     license='Apache 2.0',
-    install_requires=['imageio', 'Keras==2.2.4', 'numpy==1.16.2', 'tensorflow==1.13.1', 'tqdm'],
+    install_requires=['imageio', 'numpy', 'tensorflow==2.0.0', 'tqdm', 'pyaml'],
     extras_require={
         'tests': ['pytest==4.3.0', 'pytest-cov==2.6.1'],
         'docs': ['mkdocs==1.0.4', 'mkdocs-material==4.0.2'],
-        'gpu': ['tensorflow-gpu==1.13.1'],
+        'gpu': ['tensorflow-gpu==2.0.0'],
         'dev': ['bumpversion==0.5.3'],
     },
     classifiers=[
diff --git a/tests/models/test_models.py b/tests/models/test_models.py
index c1eac9e..7771615 100644
--- a/tests/models/test_models.py
+++ b/tests/models/test_models.py
@@ -6,7 +6,7 @@
 from ISR.models.rdn import RDN
 from ISR.models.discriminator import Discriminator
 from ISR.models.cut_vgg19 import Cut_VGG19
-from keras.optimizers import Adam
+from tensorflow.keras.optimizers import Adam
 
 
 class ModelsClassTest(unittest.TestCase):
diff --git a/tests/utils/test_metrics.py b/tests/utils/test_metrics.py
index 44e1fa5..b70a2b3 100644
--- a/tests/utils/test_metrics.py
+++ b/tests/utils/test_metrics.py
@@ -1,6 +1,6 @@
 import unittest
 import numpy as np
-import keras.backend as K
+import tensorflow.keras.backend as K
 from ISR.utils.metrics import PSNR