Make transform return the same type as input (#1182)

Author: haarisr

src/torchio/transforms/transform.py

@@ -9,6 +9,7 @@
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
+from typing import TypeVar
 from typing import Union
 
 import numpy as np
@@ -49,6 +50,8 @@
     'Superior',
 )
 
+InputType = TypeVar('InputType', bound=TypeTransformInput)
+
 
 class Transform(ABC):
     """Abstract class for all TorchIO transforms.
@@ -127,8 +130,8 @@ def __init__(
 
     def __call__(
         self,
-        data: TypeTransformInput,
-    ) -> TypeTransformInput:
+        data: InputType,
+    ) -> InputType:
         """Transform data and return a result of the same type.
 
         Args: