# Buggy code
def train(
    cls,
    images_paths: Sequence[TypePath],
    cutoff: Optional[Tuple[float, float]] = None,
    mask_path: Optional[TypePath] = None,
    masking_function: Optional[Callable] = None,
    output_path: Optional[TypePath] = None,
) -> np.ndarray:
    """Extract average histogram landmarks from images used for training.

    Args:
        images_paths: List of image paths used to train.
        cutoff: Optional minimum and maximum quantile values,
            respectively, that are used to select a range of intensity of
            interest. Equivalent to :math:`pc_1` and :math:`pc_2` in
            `Nyúl and Udupa's paper <http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.204.102&rep=rep1&type=pdf>`_.
        mask_path: Optional path to a mask image to extract voxels used for
            training.
        masking_function: Optional function used to extract voxels used for
            training.
        output_path: Optional file path with extension ``.txt`` or
            ``.npy``, where the landmarks will be saved.

    Example:

        >>> import torch
        >>> import numpy as np
        >>> from pathlib import Path
        >>> from torchio.transforms import HistogramStandardization
        >>>
        >>> t1_paths = ['subject_a_t1.nii', 'subject_b_t1.nii.gz']
        >>> t2_paths = ['subject_a_t2.nii', 'subject_b_t2.nii.gz']
        >>>
        >>> t1_landmarks_path = Path('t1_landmarks.npy')
        >>> t2_landmarks_path = Path('t2_landmarks.npy')
        >>>
        >>> t1_landmarks = (
        ...     t1_landmarks_path
        ...     if t1_landmarks_path.is_file()
        ...     else HistogramStandardization.train(t1_paths)
        ... )
        >>> torch.save(t1_landmarks, t1_landmarks_path)
        >>>
        >>> t2_landmarks = (
        ...     t2_landmarks_path
        ...     if t2_landmarks_path.is_file()
        ...     else HistogramStandardization.train(t2_paths)
        ... )
        >>> torch.save(t2_landmarks, t2_landmarks_path)
        >>>
        >>> landmarks_dict = {
        ...     't1': t1_landmarks,
        ...     't2': t2_landmarks,
        ... }
        >>>
        >>> transform = HistogramStandardization(landmarks_dict)
    """  # noqa: E501
    quantiles_cutoff = DEFAULT_CUTOFF if cutoff is None else cutoff
    percentiles_cutoff = 100 * np.array(quantiles_cutoff)
    percentiles_database = []
    percentiles = _get_percentiles(percentiles_cutoff)
    for image_file_path in tqdm(images_paths):
        tensor, _ = read_image(image_file_path)
        data = tensor.numpy()
        if masking_function is not None:
            mask = masking_function(data)
        else:
            if mask_path is not None:
                mask, _ = read_image(mask_path)
                mask = mask.numpy() > 0
            else:
                mask = np.ones_like(data, dtype=np.bool)
        percentile_values = np.percentile(data[mask], percentiles)
        percentiles_database.append(percentile_values)
    percentiles_database = np.vstack(percentiles_database)
    mapping = _get_average_mapping(percentiles_database)

    if output_path is not None:
        output_path = Path(output_path).expanduser()
        extension = output_path.suffix
        if extension == ".txt":
            modality = "image"
            text = f"{modality} {' '.join(map(str, mapping))}"
            output_path.write_text(text)
        elif extension == ".npy":
            np.save(output_path, mapping)
    return mapping


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-5-f7eeffa30bdc> in <module>()
10 ]
11 transform = tio.Compose(transforms)
---> 12 preprocessed = transform(subject)

5 frames
/usr/local/lib/python3.6/dist-packages/torchio/transforms/transform.py in __call__(self, data)
100             subject = copy.copy(subject)
101         with np.errstate(all='raise'):
--> 102             transformed = self.apply_transform(subject)
103         self.add_transform_to_subject_history(transformed)
104         for image in transformed.get_images(intensity_only=False):

/usr/local/lib/python3.6/dist-packages/torchio/transforms/augmentation/composition.py in apply_transform(self, subject)
44     def apply_transform(self, subject: Subject) -> Subject:
45         for transform in self.transforms:
---> 46             subject = transform(subject)
47         return subject
48

/usr/local/lib/python3.6/dist-packages/torchio/transforms/transform.py in __call__(self, data)
100             subject = copy.copy(subject)
101         with np.errstate(all='raise'):
--> 102             transformed = self.apply_transform(subject)
103         self.add_transform_to_subject_history(transformed)
104         for image in transformed.get_images(intensity_only=False):

/usr/local/lib/python3.6/dist-packages/torchio/transforms/preprocessing/intensity/normalization_transform.py in apply_transform(self, subject)
46     def apply_transform(self, subject: Subject) -> Subject:
47         for image_name, image in self.get_images_dict(subject).items():
---> 48             mask = Transform.get_mask(self.masking_method, subject, image.data)
49             self.apply_normalization(subject, image_name, mask)
50         return subject

/usr/local/lib/python3.6/dist-packages/torchio/transforms/transform.py in get_mask(masking_method, subject, tensor)
394             return Transform.ones(tensor)
395         elif callable(masking_method):
--> 396             return masking_method(tensor)
397         elif type(masking_method) is str:
398             if masking_method in subject and isinstance(subject[masking_method], LabelMap):

/usr/local/lib/python3.6/dist-packages/torchio/transforms/transform.py in mean(tensor)
386     @staticmethod
387     def mean(tensor: torch.Tensor) -> torch.Tensor:
--> 388         mask = tensor > tensor.mean()
389         return mask
390

RuntimeError: Can only calculate the mean of floating types. Got Short instead.

# Fixed code
def train(
    cls,
    images_paths: Sequence[TypePath],
    cutoff: Optional[Tuple[float, float]] = None,
    mask_path: Optional[TypePath] = None,
    masking_function: Optional[Callable] = None,
    output_path: Optional[TypePath] = None,
) -> np.ndarray:
    """Extract average histogram landmarks from images used for training.

    Args:
        images_paths: List of image paths used to train.
        cutoff: Optional minimum and maximum quantile values,
            respectively, that are used to select a range of intensity of
            interest. Equivalent to :math:`pc_1` and :math:`pc_2` in
            `Nyúl and Udupa's paper <http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.204.102&rep=rep1&type=pdf>`_.
        mask_path: Optional path to a mask image to extract voxels used for
            training.
        masking_function: Optional function used to extract voxels used for
            training.
        output_path: Optional file path with extension ``.txt`` or
            ``.npy``, where the landmarks will be saved.

    Example:

        >>> import torch
        >>> import numpy as np
        >>> from pathlib import Path
        >>> from torchio.transforms import HistogramStandardization
        >>>
        >>> t1_paths = ['subject_a_t1.nii', 'subject_b_t1.nii.gz']
        >>> t2_paths = ['subject_a_t2.nii', 'subject_b_t2.nii.gz']
        >>>
        >>> t1_landmarks_path = Path('t1_landmarks.npy')
        >>> t2_landmarks_path = Path('t2_landmarks.npy')
        >>>
        >>> t1_landmarks = (
        ...     t1_landmarks_path
        ...     if t1_landmarks_path.is_file()
        ...     else HistogramStandardization.train(t1_paths)
        ... )
        >>> torch.save(t1_landmarks, t1_landmarks_path)
        >>>
        >>> t2_landmarks = (
        ...     t2_landmarks_path
        ...     if t2_landmarks_path.is_file()
        ...     else HistogramStandardization.train(t2_paths)
        ... )
        >>> torch.save(t2_landmarks, t2_landmarks_path)
        >>>
        >>> landmarks_dict = {
        ...     't1': t1_landmarks,
        ...     't2': t2_landmarks,
        ... }
        >>>
        >>> transform = HistogramStandardization(landmarks_dict)
    """  # noqa: E501
    quantiles_cutoff = DEFAULT_CUTOFF if cutoff is None else cutoff
    percentiles_cutoff = 100 * np.array(quantiles_cutoff)
    percentiles_database = []
    percentiles = _get_percentiles(percentiles_cutoff)
    for image_file_path in tqdm(images_paths):
        tensor, _ = read_image(image_file_path)
        if masking_function is not None:
            mask = masking_function(tensor)
        else:
            if mask_path is not None:
                mask, _ = read_image(mask_path)
                mask = mask.numpy() > 0
            else:
                mask = np.ones_like(tensor, dtype=np.bool)
        array = tensor.numpy()
        percentile_values = np.percentile(array[mask], percentiles)
        percentiles_database.append(percentile_values)
    percentiles_database = np.vstack(percentiles_database)
    mapping = _get_average_mapping(percentiles_database)

    if output_path is not None:
        output_path = Path(output_path).expanduser()
        extension = output_path.suffix
        if extension == ".txt":
            modality = "image"
            text = f"{modality} {' '.join(map(str, mapping))}"
            output_path.write_text(text)
        elif extension == ".npy":
            np.save(output_path, mapping)
    return mapping