diff --git a/Makefile b/Makefile
index 036e9f3..5cabea2 100644
--- a/Makefile
+++ b/Makefile
@@ -1,5 +1,9 @@
 lint:
 	pre-commit run --all-files
 
+
 test:
-	python -m unittest -f tests/test_*.py
+	@for f in $(shell ls test_resnet_rs/test*.py); do \
+  		echo $${f};\
+		python $${f};\
+		done
diff --git a/scripts/test_local_output_consistency.py b/scripts/test_local_output_consistency.py
index a7b5884..09a2103 100644
--- a/scripts/test_local_output_consistency.py
+++ b/scripts/test_local_output_consistency.py
@@ -25,7 +25,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-50-i160.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs50_i160_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs50_i160_original_logits_ema.npy",
         ),
     },
     {
@@ -34,7 +35,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-101-i160.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs101_i160_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs101_i160_original_logits_ema.npy",
         ),
     },
     {
@@ -43,7 +45,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-101-i192.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs101_i192_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs101_i192_original_logits_ema.npy",
         ),
     },
     {
@@ -52,7 +55,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-152-i192.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs152_i192_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs152_i192_original_logits_ema.npy",
         ),
     },
     {
@@ -61,7 +65,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-152-i224.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs152_i224_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs152_i224_original_logits_ema.npy",
         ),
     },
     {
@@ -70,7 +75,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-152-i256.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs152_i256_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs152_i256_original_logits_ema.npy",
         ),
     },
     {
@@ -79,7 +85,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-200-i256.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs200_i256_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs200_i256_original_logits_ema.npy",
         ),
     },
     {
@@ -88,7 +95,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-270-i256.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs270_i256_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs270_i256_original_logits_ema.npy",
         ),
     },
     {
@@ -97,7 +105,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-350-i256.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs350_i256_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs350_i256_original_logits_ema.npy",
         ),
     },
     {
@@ -106,7 +115,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-350-i320.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs350_i320_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs350_i320_original_logits_ema.npy",
         ),
     },
     {
@@ -115,7 +125,8 @@
         "weights_path": os.path.join(ROOT_DIR, "weights/resnet-rs-420-i320.h5"),
         "original_outputs": os.path.join(
             ROOT_DIR,
-            "tests/assets/original_outputs/resnetrs420_i320_original_logits_ema.npy",
+            "test_resnet_rs/assets/original_outputs/"
+            "resnetrs420_i320_original_logits_ema.npy",
         ),
     },
 ]
@@ -180,7 +191,7 @@
 
 
 class TestLocalOutputConsistency(parameterized.TestCase):
-    IMAGE_PATH = os.path.join(ROOT_DIR, "tests/assets/panda.jpg")
+    IMAGE_PATH = os.path.join(ROOT_DIR, "test_resnet_rs/assets/panda.jpg")
     INPUT_SHAPE = (224, 224)
     CROP_PADDING = 32
     MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255]
diff --git a/setup.py b/setup.py
index 6137c97..7ddc7c9 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@
 from setuptools import setup
 
 TF_MIN_VERSION = 2.2
-TF_MAX_VERSION = 2.6
+TF_MAX_VERSION = 2.7
 
 with open("README.md", encoding="utf-8") as f:
     long_description = "\n" + f.read()
diff --git a/tests/assets/original_outputs/resnetrs101_i160_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs101_i160_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs101_i160_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs101_i160_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs101_i192_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs101_i192_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs101_i192_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs101_i192_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs152_i192_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs152_i192_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs152_i192_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs152_i192_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs152_i224_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs152_i224_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs152_i224_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs152_i224_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs152_i256_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs152_i256_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs152_i256_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs152_i256_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs200_i256_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs200_i256_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs200_i256_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs200_i256_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs270_i256_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs270_i256_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs270_i256_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs270_i256_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs350_i256_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs350_i256_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs350_i256_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs350_i256_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs350_i320_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs350_i320_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs350_i320_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs350_i320_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs420_i320_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs420_i320_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs420_i320_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs420_i320_original_logits_ema.npy
diff --git a/tests/assets/original_outputs/resnetrs50_i160_original_logits_ema.npy b/test_resnet_rs/assets/original_outputs/resnetrs50_i160_original_logits_ema.npy
similarity index 100%
rename from tests/assets/original_outputs/resnetrs50_i160_original_logits_ema.npy
rename to test_resnet_rs/assets/original_outputs/resnetrs50_i160_original_logits_ema.npy
diff --git a/tests/assets/panda.jpg b/test_resnet_rs/assets/panda.jpg
similarity index 100%
rename from tests/assets/panda.jpg
rename to test_resnet_rs/assets/panda.jpg
diff --git a/tests/test_model_utils.py b/test_resnet_rs/test_model_utils.py
similarity index 100%
rename from tests/test_model_utils.py
rename to test_resnet_rs/test_model_utils.py
diff --git a/tests/test_models.py b/test_resnet_rs/test_models.py
similarity index 100%
rename from tests/test_models.py
rename to test_resnet_rs/test_models.py