diff --git a/tests/hparams_optim/test_mlflow_with_optuna.py b/examples/example_mlflow_optuna_cifar10_demo.py
similarity index 77%
rename from tests/hparams_optim/test_mlflow_with_optuna.py
rename to examples/example_mlflow_optuna_cifar10_demo.py
index 317378f..ebdbcc1 100644
--- a/tests/hparams_optim/test_mlflow_with_optuna.py
+++ b/examples/example_mlflow_optuna_cifar10_demo.py
@@ -1,5 +1,11 @@
-''' Script created by PeterC to test mlflowand optuna integration library for model monitoring and tracking - 02-07-2024 '''
+"""MLflow and Optuna CIFAR10 demo script.
 
+This is intentionally an example, not a pytest module: it downloads CIFAR10,
+starts long optimization runs, and may start external tracking services.
+"""
+
+import numpy as np
+import pyTorchAutoForge  # Custom torch tools
 import optuna
 import mlflow
 import mlflow.pytorch
@@ -11,24 +17,25 @@
 import matplotlib.pyplot as plt
 
 # Import modules
-import os, subprocess, time, logging
+import os
+import subprocess
+import time
+import logging
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-import pyTorchAutoForge # Custom torch tools
-import numpy as np
 
-def StartMLflowUI(port:int=5000):
+def StartMLflowUI(port: int = 5000):
 
     # Start MLflow UI
     os.system('mlflow ui --port ' + str(port))
-    process = subprocess.Popen(['mlflow', 'ui', '--port ' + f'{port}', '&'], 
+    process = subprocess.Popen(['mlflow', 'ui', '--port ' + f'{port}', '&'],
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     print(f'MLflow UI started with PID: {process.pid}, on port: {port}')
-    time.sleep(1) # Ensure the server has started
+    time.sleep(1)  # Ensure the server has started
     if process.poll() is None:
         print('MLflow UI is running OK.')
     else:
@@ -36,16 +43,22 @@ def StartMLflowUI(port:int=5000):
 
     return process
 
+
 # %% Datasets loading (global)
 # Load CIFAR-10 dataset from torchvision
-transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
-train_dataset = datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
-test_dataset = datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
+transform = transforms.Compose(
+    [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+train_dataset = datasets.CIFAR10(
+    root='./data', train=True, download=True, transform=transform)
+test_dataset = datasets.CIFAR10(
+    root='./data', train=False, download=True, transform=transform)
 train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
 test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
 
 # TEST EXAMPLE BY GPT
 # %% Optuna model for trial optimization
+
+
 class SimpleCNN(nn.Module):
     def __init__(self, trial):
         super(SimpleCNN, self).__init__()
@@ -56,11 +69,11 @@ def __init__(self, trial):
 
         # NOTE: What are the second inputs to suggest_int? Are they the lower and upper bounds?
 
-        in_channels = 3 # Number of channels in the input image (RGB)
+        in_channels = 3  # Number of channels in the input image (RGB)
 
         # Add convolutional blocks to the model up to "num_conv_layers" --> optimization parameter
         for i in range(num_conv_layers):
-            
+
             # Number of output channels for the ith convolutional block
             out_channels = trial.suggest_int(f'filters_{i}', 16, 64)
             # Kernel size for the ith convolutional block
@@ -68,13 +81,16 @@ def __init__(self, trial):
 
             # Add convolutional block to the model. layers.add_module() takes input:(name, module)
             # NOTE: this is easily repurposed for the Model AutoBuilder
-            self.layers.add_module(f'conv{i}', nn.Conv2d(in_channels, out_channels, kernel_size))
+            self.layers.add_module(f'conv{i}', nn.Conv2d(
+                in_channels, out_channels, kernel_size))
             self.layers.add_module(f'relu{i}', nn.ReLU())
             self.layers.add_module(f'maxpool{i}', nn.MaxPool2d(2))
 
-            in_channels = out_channels # Number of input channels to conv2d is the number of output channels from the previous conv2d
+            # Number of input channels to conv2d is the number of output channels from the previous conv2d
+            in_channels = out_channels
 
-        self.layers.add_module('flatten', nn.Flatten()) # Add flatten before NN
+        self.layers.add_module('flatten', nn.Flatten()
+                               )  # Add flatten before NN
 
         # Dense layers
         num_dense_layers = trial.suggest_int('num_dense_layers', 1, 3)
@@ -82,7 +98,8 @@ def __init__(self, trial):
 
         for i in range(num_dense_layers):
             out_features = trial.suggest_int(f'units_{i}', 32, 256)
-            self.layers.add_module(f'fc{i}', nn.Linear(in_features, out_features))
+            self.layers.add_module(
+                f'fc{i}', nn.Linear(in_features, out_features))
             self.layers.add_module(f'relu_fc{i}', nn.ReLU())
             dropout_rate = trial.suggest_float(f'dropout_rate_{i}', 0.1, 0.5)
             self.layers.add_module(f'dropout{i}', nn.Dropout(dropout_rate))
@@ -101,6 +118,8 @@ def _get_conv_output(self, shape):
             return output.size(1)
 
 # %% Optuna objective function
+
+
 def objective(trial):
     # Create new run_IN within the current session
     with mlflow.start_run():
@@ -111,11 +130,13 @@ def objective(trial):
         model = SimpleCNN(trial).to(device)
 
         # Select the optimizer
-        optimizer_name = trial.suggest_categorical('optimizer', ['Adam', 'RMSprop'])
-        lr = trial.suggest_float('lr', 1e-6, 1e-2, log=True) # NOTE: What are the inputs to suggest_float?
+        optimizer_name = trial.suggest_categorical(
+            'optimizer', ['Adam', 'RMSprop'])
+        # NOTE: What are the inputs to suggest_float?
+        lr = trial.suggest_float('lr', 1e-6, 1e-2, log=True)
         optimizer = getattr(optim, optimizer_name)(model.parameters(), lr=lr)
         # NOTE: What is "getattr" function? It should be a pytorch function: likely a method of "model", TBC
-    
+
         # Loss function definition
         criterion = nn.CrossEntropyLoss()
 
@@ -137,12 +158,12 @@ def objective(trial):
             mlflow.log_param(f'units_{i}', trial.params.get(f'units_{i}'))
             mlflow.log_param(
                 f'dropout_rate_{i}', trial.params.get(f'dropout_rate_{i}'))
-            
+
         # Training the current trial model
         for epoch in range(10):
             model.train()
             for batch in train_loader:
-                inputs, targets = batch # Unpack tuples
+                inputs, targets = batch  # Unpack tuples
                 inputs, targets = inputs.to(device), targets.to(device)
                 optimizer.zero_grad()
 
@@ -168,19 +189,21 @@ def objective(trial):
                     correct += (predicted == targets).sum().item()
 
             accuracy = correct / total
-            
+
             # Log the accuracy of the model using mlflow
             mlflow.log_metric('Accuracy value', accuracy, step=epoch)
 
-            trial.report(accuracy, epoch) # Report the accuracy of the model to optuna pruner
+            # Report the accuracy of the model to optuna pruner
+            trial.report(accuracy, epoch)
 
             if trial.should_prune():
                 # Mark the run as killed in mlflow
                 mlflow.end_run(status='KILLED')
-                raise optuna.TrialPruned() # Raise an optuna exception to stop the trial due to pruning
-            
+                # Raise an optuna exception to stop the trial due to pruning
+                raise optuna.TrialPruned()
+
         # MLflow: log trial data (it will have a unique ID)
-        # NOTE: this could be matched with the AutoBuilder configuration such that the parameters 
+        # NOTE: this could be matched with the AutoBuilder configuration such that the parameters
         # of config file are automatically logged by mlflow with their names
         mlflow.end_run(status='FINISHED')
 
@@ -191,42 +214,45 @@ def main():
     print('---------------------------- TEST SCRIPT: MLflow and Optuna functionalities ----------------------------\n')
     # %% MLflow tracking initialization
     port = 7000
-    #StartMLflowUI(port) # Start MLflow UI
+    # StartMLflowUI(port) # Start MLflow UI
 
     # TODO: add change of script execution working directory to the current directory
-    
+
     # %% Optuna study configuration
     if not (os.path.exists('testdata/optuna_db')):
         os.makedirs('testdata/optuna_db')
 
     studyName = 'CIFAR10_CNN_OptimizationExample'
     optunaStudyObj = optuna.create_study(study_name=studyName,
-                                         storage='sqlite:///testdata/{studyName}.db'.format(studyName=os.path.join('optuna_db', studyName)),
+                                         storage='sqlite:///testdata/{studyName}.db'.format(
+                                             studyName=os.path.join('optuna_db', studyName)),
                                          load_if_exists=True,
                                          direction='maximize',
                                          sampler=optuna.samplers.TPESampler(),
                                          pruner=optuna.pruners.SuccessiveHalvingPruner(min_resource=1, reduction_factor=2,
                                                                                        min_early_stopping_rate=1))
-    
+
     # Mlflow experiment name
     mlflow.set_experiment(studyName)
 
     # %% Optuna optimization
     optunaStudyObj.optimize(objective, n_trials=100, timeout=1800)
-    
+
     # Print the best trial
-    print('Number of finished trials:', len(optunaStudyObj.trials)) # Get number of finished trials
+    # Get number of finished trials
+    print('Number of finished trials:', len(optunaStudyObj.trials))
     print('Best trial:')
 
-    trial = optunaStudyObj.best_trial # Get the best trial from study object
-    print('  Value: {:.4f}'.format(trial.value)) # Loss function value for the best trial
+    trial = optunaStudyObj.best_trial  # Get the best trial from study object
+    # Loss function value for the best trial
+    print('  Value: {:.4f}'.format(trial.value))
 
     # Print parameters of the best trial
     print('  Params: ')
     for key, value in trial.params.items():
         print('    {}: {}'.format(key, value))
-    
+
 
 # %% Main call
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/pyTorchAutoForge/api/torch/torchModulesIO.py b/pyTorchAutoForge/api/torch/torchModulesIO.py
index a5ab6ac..8cbe886 100644
--- a/pyTorchAutoForge/api/torch/torchModulesIO.py
+++ b/pyTorchAutoForge/api/torch/torchModulesIO.py
@@ -36,7 +36,7 @@ def SaveModel(model: torch.nn.Module,
               model_filename: str | pathlib.Path, 
               save_mode : AutoForgeModuleSaveMode | str = AutoForgeModuleSaveMode.MODEL_ARCH_STATE, 
               example_input: torch.Tensor | None = None, 
-              target_device: str = 'cpu', 
+              target_device: str | torch.device = 'cpu',
               model_base_name : str | None = None) -> None:
     """
     Saves a PyTorch model to a file.
@@ -57,7 +57,7 @@ def SaveModel(model: torch.nn.Module,
             Defaults to AutoForgeModuleSaveMode.MODEL_ARCH_STATE.
         example_input (torch.Tensor | None, optional): A sample input tensor for tracing or scripting.
             Defaults to None.
-        target_device (str, optional): The target device (e.g., 'cpu' or 'cuda:0') to save the model.
+        target_device (str | torch.device, optional): The target device (e.g., 'cpu' or 'cuda:0') to save the model.
             Defaults to 'cpu'.
         model_base_name (str | None, optional): An optional base name for the model.
             Defaults to None.
@@ -94,9 +94,10 @@ def SaveModel(model: torch.nn.Module,
     else: 
         extension = '.pth'
         
+    target_device = torch.device(target_device)
+
     # Format target device string to remove ':' from name
-    target_device_name = target_device
-    target_device_name = target_device_name.replace(':', '')
+    target_device_name = str(target_device).replace(':', '')
 
     # Form filename for saving
     # Check if device is in model name and remove it
@@ -283,5 +284,3 @@ def ValidateDictLoading(model: torch.nn.Module | torch.nn.ModuleDict | torch.nn.
 
     else:
         print("All model parameters are correctly loaded.")
-
-
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..a18217b
--- /dev/null
+++ b/tests/__init__.py
@@ -0,0 +1 @@
+"""Shared pytest support package for local tests."""
diff --git a/tests/api/tcp/test_start_server.py b/tests/api/tcp/test_start_server.py
index 00756bf..b59da74 100644
--- a/tests/api/tcp/test_start_server.py
+++ b/tests/api/tcp/test_start_server.py
@@ -1,50 +1,64 @@
-import numpy as np
+from __future__ import annotations
 
-# Custom imports
-from pyTorchAutoForge.api.tcp import DataProcessor, pytcp_server, pytcp_requestHandler, ProcessingMode
 import threading
 
-# MAIN SCRIPT
-def main():
-    print('\n\n----------------------------------- RUNNING: test_start_server.py -----------------------------------\n')
-    
-    # %% TCP SERVER INITIALIZATION
-    HOST1, PORT1 = "127.0.0.1", 50000 # Define host and port for the first server
-    HOST2, PORT2 = "127.0.0.1", 50001 # Define host and port for the second server
-
-    def dummy_function(inputData):
-        return inputData
-
-    # Define DataProcessor object for RequestHandler
-    dataProcessorObj = DataProcessor(dummy_function, np.float32, 1024,
-                                        ENDIANNESS='little', DYNAMIC_BUFFER_MODE=True,
-                                        PRE_PROCESSING_MODE=ProcessingMode.TENSOR)
-    
-    dataProcessorObj_multi = DataProcessor(dummy_function, np.float32, 1024, 
-                                                 ENDIANNESS='little', DYNAMIC_BUFFER_MODE=True, 
-                                                 PRE_PROCESSING_MODE=ProcessingMode.MULTI_TENSOR)
-
-    def start_server(host, port, dataProcessorObj):
-        with pytcp_server((host, port), pytcp_requestHandler, dataProcessorObj, bindAndActivate=True) as server:
-            try:
-                print(f'\nServer initialized correctly on {host}:{port}. Set in "serve_forever" mode.')
-                server.serve_forever()
-            except KeyboardInterrupt:
-                print(f"\nServer on {host}:{port} is gracefully shutting down =D.")
-                server.shutdown()
-                server.server_close()
-
-    # Start two servers on separate threads
-    thread1 = threading.Thread(target=start_server, args=( HOST1, PORT1, dataProcessorObj) )
-    thread2 = threading.Thread(target=start_server, args=( HOST2, PORT2, dataProcessorObj_multi) )
-
-    thread1.start()
-    thread2.start()
-
-    # Wait for the threads to finish processing
-    thread1.join()
-    thread2.join()
-
-if __name__ == "__main__":
-    main()
+import numpy as np
+
+from pyTorchAutoForge.api.tcp import (
+    DataProcessor,
+    ProcessingMode,
+    pytcp_requestHandler,
+    pytcp_server,
+)
+
+
+def _Identity(input_data_: np.ndarray) -> np.ndarray:
+    return input_data_
+
+
+def test_data_processor_tensor_roundtrip() -> None:
+    processor_ = DataProcessor(
+        _Identity,
+        np.float32,
+        1024,
+        ENDIANNESS="little",
+        DYNAMIC_BUFFER_MODE=True,
+        PRE_PROCESSING_MODE=ProcessingMode.TENSOR,
+    )
+    input_array_ = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32)
+
+    output_buffer_ = processor_.process(
+        processor_.TensorToBytesBuffer(input_array_))
+    output_array_, output_shape_ = processor_.BytesBufferToTensor(output_buffer_[
+                                                                  4:])
+
+    assert output_shape_ == input_array_.shape
+    assert np.array_equal(output_array_, input_array_)
+
+
+def test_tcp_server_starts_and_shuts_down_without_hanging() -> None:
+    processor_ = DataProcessor(
+        _Identity,
+        np.float32,
+        1024,
+        ENDIANNESS="little",
+        DYNAMIC_BUFFER_MODE=True,
+        PRE_PROCESSING_MODE=ProcessingMode.TENSOR,
+    )
+
+    with pytcp_server(
+        ("127.0.0.1", 0),
+        pytcp_requestHandler,
+        processor_,
+        bindAndActivate=True,
+    ) as server_:
+        thread_ = threading.Thread(target=server_.serve_forever, daemon=True)
+        thread_.start()
+
+        assert server_.server_address[0] == "127.0.0.1"
+        assert server_.server_address[1] > 0
+
+        server_.shutdown()
+        thread_.join(timeout=2.0)
 
+    assert not thread_.is_alive()
diff --git a/tests/api/tcp/test_torch_model_over_tcp.py b/tests/api/tcp/test_torch_model_over_tcp.py
deleted file mode 100644
index 57639ba..0000000
--- a/tests/api/tcp/test_torch_model_over_tcp.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import pyTorchAutoForge
-
-def main():
-    pass
-
-
-if __name__ == "__main__":
-    main()
\ No newline at end of file
diff --git a/tests/api/torch/test_torch_modules_io.py b/tests/api/torch/test_torch_modules_io.py
new file mode 100644
index 0000000..2cb125c
--- /dev/null
+++ b/tests/api/torch/test_torch_modules_io.py
@@ -0,0 +1,19 @@
+from pathlib import Path
+
+import torch
+
+from pyTorchAutoForge.api.torch import AutoForgeModuleSaveMode, SaveModel
+
+
+def test_SaveModel_accepts_torch_device(tmp_path: Path) -> None:
+    model = torch.nn.Linear(2, 1)
+    model_path = tmp_path / "linear_model"
+
+    SaveModel(
+        model=model,
+        model_filename=model_path,
+        save_mode=AutoForgeModuleSaveMode.MODEL_STATE_DICT,
+        target_device=torch.device("cpu"),
+    )
+
+    assert (tmp_path / "linear_model_statedict.pth").is_file()
diff --git a/tests/conftest.py b/tests/conftest.py
index 6fa9f72..063936e 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -19,6 +19,43 @@
 )
 
 
+def pytest_addoption(parser: pytest.Parser) -> None:
+    parser.addoption(
+        "--run-slow",
+        action="store_true",
+        default=False,
+        help="Run tests marked slow.",
+    )
+    parser.addoption(
+        "--run-gpu",
+        action="store_true",
+        default=False,
+        help="Run tests marked gpu.",
+    )
+    parser.addoption(
+        "--run-visual",
+        action="store_true",
+        default=False,
+        help="Run tests marked visual.",
+    )
+
+
+def pytest_collection_modifyitems(config: pytest.Config,
+                                  items: list[pytest.Item],
+                                  ) -> None:
+    marker_options_ = {
+        "slow": ("--run-slow", "slow test skipped by default"),
+        "gpu": ("--run-gpu", "GPU test skipped by default"),
+        "visual": ("--run-visual", "visual test skipped by default"),
+    }
+
+    for item_ in items:
+        for marker_name_, (option_name_, reason_) in marker_options_.items():
+            if marker_name_ in item_.keywords and not config.getoption(option_name_):
+                item_.add_marker(pytest.mark.skip(reason=reason_))
+                break
+
+
 def _cleanup_legacy_output_dirs() -> None:
     repo_root_resolved_ = REPO_ROOT.resolve()
     for output_dir_ in LEGACY_TEST_OUTPUT_DIRS:
diff --git a/tests/datasets/test_ImagesAugmentation.py b/tests/datasets/test_ImagesAugmentation.py
index c57c04b..40741ac 100644
--- a/tests/datasets/test_ImagesAugmentation.py
+++ b/tests/datasets/test_ImagesAugmentation.py
@@ -10,6 +10,7 @@
 import matplotlib.pyplot as plt
 from pyTorchAutoForge.datasets import ImagesLabelsCachedDataset, AugmentationConfig, ImageAugmentationsHelper
 from pyTorchAutoForge.utils.conversion_utils import torch_to_numpy, numpy_to_torch
+from tests.helpers import CudaIsUsable, RequireCudaUsable
 import numpy as np
 from time import perf_counter
 import PIL
@@ -194,6 +195,9 @@ def test_detect_border_crossing_white_masks():
 # %% Integrated tests
 
 
+@pytest.mark.integration
+@pytest.mark.slow
+@pytest.mark.visual
 def test_synthetic_mask_augmentation():
 
     augs_datakey = [DataKey.IMAGE, DataKey.KEYPOINTS]
@@ -219,7 +223,7 @@ def test_synthetic_mask_augmentation():
         contrast_aug_prob=1.0,
         input_normalization_factor=255.0,
         enable_auto_input_normalization=True,
-        device='cuda' if torch.cuda.is_available() else 'cpu'
+        device='cuda' if CudaIsUsable() else 'cpu'
     )
 
     augs_helper = ImageAugmentationsHelper(cfg)
@@ -281,6 +285,9 @@ def test_synthetic_mask_augmentation():
     plt.close()
 
 
+@pytest.mark.integration
+@pytest.mark.slow
+@pytest.mark.visual
 def test_sample_images_augmentation():
     # Load sample images
     imgs = _load_sample_images()
@@ -330,7 +337,7 @@ def test_sample_images_augmentation():
                              input_normalization_factor=255.0,
                              enable_auto_input_normalization=True,
                              enable_batch_validation_check=True,
-                             device='cuda' if torch.cuda.is_available() else 'cpu'
+                             device='cuda' if CudaIsUsable() else 'cpu'
                              )
 
     augs_helper = ImageAugmentationsHelper(cfg)
@@ -380,6 +387,7 @@ def test_sample_images_augmentation():
         plt.close()
 
 
+@pytest.mark.visual
 def test_random_softbinarize_only_sample_images_with_viz():
     # Load sample images
     np.random.seed(0)
@@ -455,6 +463,8 @@ def test_random_softbinarize_only_sample_images_with_viz():
     plt.close()
 
 
+@pytest.mark.integration
+@pytest.mark.visual
 def test_AugmentationSequential():
     from kornia.augmentation import AugmentationSequential, RandomAffine, RandomHorizontalFlip
     from kornia.constants import DataKey
@@ -542,7 +552,7 @@ def augment_data_batch(*inputs: torch.Tensor) -> tuple[torch.Tensor, ...]:
 # Chain parametrize to test combinations
 
 
-@pytest.mark.parametrize("device", ["cpu", "cuda"])
+@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)])
 @pytest.mark.parametrize("shift_aug_prob", [0, 1])
 @pytest.mark.parametrize("rotation_aug_prob", [0, 1])
 @pytest.mark.parametrize("gaussian_noise_aug_prob", [0, 1])
@@ -557,8 +567,8 @@ def test_augmentation_helper_preserves_device(device,
                                               brightness_aug_prob,
                                               contrast_aug_prob) -> None:
 
-    if device == "cuda" and not torch.cuda.is_available():
-        pytest.skip("CUDA not available.")
+    if device == "cuda":
+        RequireCudaUsable()
 
     # Create a minimal configuration for the augmentations helper.
     cfg = AugmentationConfig(
@@ -601,10 +611,11 @@ def test_augmentation_helper_preserves_device(device,
     _assert_helper_modules_on_device(augs_helper, device_t)
 
 
+@pytest.mark.gpu
+@pytest.mark.slow
 def test_augmentation_helper_timing_cpu_vs_cuda():
     """Test execution time cpu vs cuda for augmentations helper"""
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA not available.")
+    RequireCudaUsable()
     num_iters_ = 500
     cfg_kwargs = dict(
         max_shift_img_fraction=(0.2, 0.2),
@@ -1106,7 +1117,7 @@ def test_helper_metadata_supports_external_sun_direction_label_update():
 # %% MANUAL TEST CALLS
 if __name__ == '__main__':
 
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cuda" if CudaIsUsable() else "cpu"
     shift_aug_prob = 0.0
     rotation_aug_prob = 0.0
     gaussian_noise_aug_prob = 0.0
diff --git a/tests/datasets/test_auxiliary_functions.py b/tests/datasets/test_auxiliary_functions.py
index 837eda8..1290d6e 100644
--- a/tests/datasets/test_auxiliary_functions.py
+++ b/tests/datasets/test_auxiliary_functions.py
@@ -3,9 +3,14 @@
 
 import os
 import numpy as np
+import pytest
 
 from pyTorchAutoForge.setup.AutoForgeInit import Is_session_headless
 
+
+@pytest.mark.integration
+@pytest.mark.slow
+@pytest.mark.visual
 def test_LoadDatasetToMem():
 
     import matplotlib
diff --git a/tests/evaluation/test_ModelExplainer.py b/tests/evaluation/test_ModelExplainer.py
deleted file mode 100644
index f87f5c1..0000000
--- a/tests/evaluation/test_ModelExplainer.py
+++ /dev/null
@@ -1 +0,0 @@
-# TODO
\ No newline at end of file
diff --git a/tests/helpers.py b/tests/helpers.py
new file mode 100644
index 0000000..183b671
--- /dev/null
+++ b/tests/helpers.py
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+import pytest
+import torch
+
+
+@lru_cache(maxsize=1)
+def CudaIsUsable() -> bool:
+    """Return True only when CUDA kernels can execute, not just be discovered."""
+    if not torch.cuda.is_available():
+        return False
+
+    try:
+        device_ = torch.device("cuda")
+        tensor_ = torch.ones(1, device=device_)
+        result_ = (tensor_ + 1.0).item()
+        torch.cuda.synchronize(device_)
+        return result_ == 2.0
+    except Exception:
+        return False
+
+
+def RequireCudaUsable() -> None:
+    if not CudaIsUsable():
+        pytest.skip("CUDA runtime is not usable in this environment.")
diff --git a/tests/model_building/test_TorchModel.py b/tests/model_building/test_TorchModel.py
deleted file mode 100644
index c5fe3d3..0000000
--- a/tests/model_building/test_TorchModel.py
+++ /dev/null
@@ -1,13 +0,0 @@
-import torch 
-from pyTorchAutoForge.utils import GetDevice
-from pyTorchAutoForge.api.torch.torchModulesIO import *
-from pyTorchAutoForge.model_building import AutoForgeModule
-
-
-
-def main():
-    pass
-
-if __name__ == '__main__':
-    main()
-
diff --git a/tests/model_building/test_image_processing_operators.py b/tests/model_building/test_image_processing_operators.py
index f42ef8f..2dd856f 100644
--- a/tests/model_building/test_image_processing_operators.py
+++ b/tests/model_building/test_image_processing_operators.py
@@ -5,6 +5,7 @@
 from pyTorchAutoForge.datasets.LabelsClasses import PTAF_Datakey
 from kornia.constants import DataKey
 import torch
+import pytest
 from pyTorchAutoForge.model_building.backbones.image_processing_operators import Compute_threshold_mask, Apply_sobel_gradient, Apply_laplacian_of_gaussian, Compute_distance_transform_map, Compute_local_variance_map
 from pyTorchAutoForge.utils import torch_to_numpy, numpy_to_torch
 from PIL import Image
@@ -192,6 +193,8 @@ def _run_all_batched_images_(image_names, apply_augs, augmentation_module: Image
     plt.close()
 
 
+@pytest.mark.integration
+@pytest.mark.visual
 def test_all_operators():
     # ---- Configuration ----
     this_file_path = os.path.dirname(os.path.abspath(__file__))
diff --git a/tests/model_building/test_models_export.py b/tests/model_building/test_models_export.py
index bd18df6..87ec769 100644
--- a/tests/model_building/test_models_export.py
+++ b/tests/model_building/test_models_export.py
@@ -1,8 +1,10 @@
 import torch
+import pytest
 from pathlib import Path
 from pyTorchAutoForge.model_building.backbones.efficient_net import EfficientNetConfig, FeatureExtractorFactory
 
 # Export efficient net to ONNX
+@pytest.mark.export
 def test_efficientnet_basic_backbone_export(tmp_path: Path):
     # Create configuration
     cfg = EfficientNetConfig(
diff --git a/tests/optimization/test_model_training_manager.py b/tests/optimization/test_model_training_manager.py
index 6f31e7f..431b401 100644
--- a/tests/optimization/test_model_training_manager.py
+++ b/tests/optimization/test_model_training_manager.py
@@ -7,8 +7,13 @@
 
 from torchvision import models
 import optuna, mlflow
+from tests.helpers import CudaIsUsable
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA GPU required for this test.")
+
+@pytest.mark.gpu
+@pytest.mark.integration
+@pytest.mark.slow
+@pytest.mark.skipif(not CudaIsUsable(), reason="Usable CUDA runtime required for this test.")
 def test_ModelTrainingManager():
 
     from torchvision import transforms
diff --git a/tests/utils/test_timing_utils.py b/tests/utils/test_timing_utils.py
index 854edc3..3af9c98 100644
--- a/tests/utils/test_timing_utils.py
+++ b/tests/utils/test_timing_utils.py
@@ -1,40 +1,63 @@
+from __future__ import annotations
+
+from itertools import count
+
+import pytest
+
+import pyTorchAutoForge.utils.timing_utils as timing_utils
 from pyTorchAutoForge.utils import timeit_averaged, timeit_averaged_
-import time
 
-@timeit_averaged(2)
-def dummy_function():
-    print("Dummy function called")
-    time.sleep(1)
 
-def test_timeit_averaged():
-    print("Testing timeit_averaged as decorator...")
-    dummy_function()
+def _PatchPerfCounter(monkeypatch: pytest.MonkeyPatch,
+                      step: float = 0.01,
+                      ) -> None:
+    counter_ = count()
+
+    def FakePerfCounter() -> float:
+        return next(counter_) * step
+
+    monkeypatch.setattr(timing_utils.time, "perf_counter", FakePerfCounter)
+
+
+def test_timeit_averaged_returns_wrapped_result(monkeypatch: pytest.MonkeyPatch) -> None:
+    _PatchPerfCounter(monkeypatch)
+    calls_: list[int] = []
+
+    @timeit_averaged(2)
+    def SampleFunction(value_: int) -> int:
+        calls_.append(value_)
+        return value_ + 1
+
+    assert SampleFunction(4) == 5
+    assert calls_ == [4, 4]
+
+
+def test_timeit_averaged_preserves_function_metadata(monkeypatch: pytest.MonkeyPatch) -> None:
+    _PatchPerfCounter(monkeypatch)
+
+    @timeit_averaged(1)
+    def SampleFunction() -> str:
+        return "ok"
+
+    assert SampleFunction.__name__ == "SampleFunction"
+    assert SampleFunction() == "ok"
 
-def test_timeit_averaged_wrapped():
-    print("Testing timeit_averaged wrapped...")
-    dummy_function()
 
-def test_timeit_averaged_function():
-    def sample_function(x, y):
-        time.sleep(0.5)
-        return x + y
+def test_timeit_averaged_function() -> None:
+    def SampleFunction(x_: int, y_: int) -> int:
+        return x_ + y_
 
-    num_trials = 3
-    average_time = timeit_averaged_(sample_function, num_trials, 2, 3)
+    average_time_ = timeit_averaged_(SampleFunction, 3, 2, 3)
 
-    assert isinstance(average_time, float), "Average time should be a float"
-    assert average_time > 0, "Average time should be greater than 0"
-    assert average_time < 1, "Average time should be less than 1 for this test case"
+    assert isinstance(average_time_, float)
+    assert average_time_ >= 0.0
 
-def test_timeit_averaged_function_with_kwargs():
-    def sample_function(x, y, delay=0.5):
-        time.sleep(delay)
-        return x * y
 
-    num_trials = 2
-    average_time = timeit_averaged_(sample_function, num_trials, 3, 4, delay=0.3)
+def test_timeit_averaged_function_with_kwargs() -> None:
+    def SampleFunction(x_: int, y_: int, scale_: int = 1) -> int:
+        return scale_ * x_ * y_
 
-    assert isinstance(average_time, float), "Average time should be a float"
-    assert average_time > 0, "Average time should be greater than 0"
-    assert average_time < 1, "Average time should be less than 1 for this test case"
+    average_time_ = timeit_averaged_(SampleFunction, 2, 3, 4, scale_=2)
 
+    assert isinstance(average_time_, float)
+    assert average_time_ >= 0.0