project-lighter
diff --git a/‎lighter/callbacks/utils.py
+3-2 b/‎lighter/callbacks/utils.py
+3-2
diff --git a/‎lighter/callbacks/writer/base.py
+3-3 b/‎lighter/callbacks/writer/base.py
+3-3
diff --git a/‎lighter/callbacks/writer/file.py
+3-6 b/‎lighter/callbacks/writer/file.py
+3-6
@@ -1,5 +1,6 @@
 import torch
 import torchvision
+from torch import Tensor
 
 
 def get_lighter_mode(lightning_stage: str) -> str:
@@ -15,13 +16,13 @@ def get_lighter_mode(lightning_stage: str) -> str:
     return lightning_to_lighter[lightning_stage]
 
 
-def preprocess_image(image: torch.Tensor) -> torch.Tensor:
+def preprocess_image(image: Tensor) -> Tensor:
     """Preprocess the image before logging it. If it is a batch of multiple images,
     it will create a grid image of them. In case of 3D, a single image is displayed
     with slices stacked vertically, while a batch of 3D images as a grid where each
     column is a different 3D image.
     Args:
-        image (torch.Tensor): A 2D or 3D image tensor.
+        image (Tensor): A 2D or 3D image tensor.
     Returns:
         The image ready for logging.
     """
 
@@ -2,12 +2,12 @@
 
 import gc
 from abc import ABC, abstractmethod
-from datetime import datetime
 from pathlib import Path
 
 import torch
 from loguru import logger
 from pytorch_lightning import Callback, Trainer
+from torch import Tensor
 
 from lighter import LighterSystem
 
@@ -48,7 +48,7 @@ def writers(self) -> Dict[str, Callable]:
         """
 
     @abstractmethod
-    def write(self, tensor: torch.Tensor, id: int) -> None:
+    def write(self, tensor: Tensor, id: int) -> None:
         """
         Method to define how a tensor should be saved. The input tensor will be a single tensor without
         the batch dimension.
@@ -57,7 +57,7 @@ def write(self, tensor: torch.Tensor, id: int) -> None:
         A specific writer function can be retrieved using `self.get_writer(self.format)`.
 
         Args:
-            tensor (torch.Tensor): Tensor, without the batch dimension, to be saved.
+            tensor (Tensor): Tensor, without the batch dimension, to be saved.
             id (int): Identifier for the tensor, can be used for naming files or adding table records.
         """
 
 
@@ -1,12 +1,12 @@
 from typing import Callable, Dict, Union
 
 from functools import partial
-from pathlib import Path
 
 import torch
 import torchvision
 from monai.data import metatensor_to_itk_image
 from monai.transforms import DivisiblePad
+from torch import Tensor
 
 from lighter.callbacks.utils import preprocess_image
 from lighter.callbacks.writer.base import LighterBaseWriter
@@ -27,9 +27,6 @@ class LighterFileWriter(LighterBaseWriter):
             `tensor` is a single tensor without the batch dimension.
     """
 
-    def __init__(self, path: Union[str, Path], writer: Union[str, Callable]) -> None:
-        super().__init__(path, writer)
-
     @property
     def writers(self) -> Dict[str, Callable]:
         return {
@@ -41,7 +38,7 @@ def writers(self) -> Dict[str, Callable]:
             "itk_nifti": partial(write_itk_image, suffix=".nii.gz"),
         }
 
-    def write(self, tensor: torch.Tensor, id: Union[int, str]) -> None:
+    def write(self, tensor: Tensor, id: Union[int, str]) -> None:
         """
         Write the tensor using the writer specified at the instatiation.
 
@@ -82,7 +79,7 @@ def write_video(path, tensor):
     torchvision.io.write_video(str(path), tensor, fps=24)
 
 
-def write_itk_image(path: str, tensor: torch.Tensor, suffix) -> None:
+def write_itk_image(path: str, tensor: Tensor, suffix) -> None:
     path = path.with_suffix(suffix)
     itk_image = metatensor_to_itk_image(tensor, channel_dim=0, dtype=tensor.dtype)
     OPTIONAL_IMPORTS["itk"].imwrite(itk_image, str(path), True)