braindecode 0.8.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/regressor.py

@@ -11,8 +11,8 @@ import warnings
 import numpy as np
 from skorch.regressor import NeuralNetRegressor
 
-from .training.scoring import predict_trials
 from .eegneuralnet import _EEGNeuralNet
+from .training.scoring import predict_trials
 from .util import ThrowAwayIndexLoader, update_estimator_docstring
 
 
@@ -58,19 +58,28 @@ class EEGRegressor(_EEGNeuralNet, NeuralNetRegressor):
     """  # noqa: E501
     __doc__ = update_estimator_docstring(NeuralNetRegressor, doc)
 
-    def __init__(self, module, *args, cropped=False, callbacks=None,
-                 iterator_train__shuffle=True,
-                 iterator_train__drop_last=True,
-                 aggregate_predictions=True, **kwargs):
+    def __init__(
+        self,
+        module,
+        *args,
+        cropped=False,
+        callbacks=None,
+        iterator_train__shuffle=True,
+        iterator_train__drop_last=True,
+        aggregate_predictions=True,
+        **kwargs,
+    ):
         self.cropped = cropped
         self.aggregate_predictions = aggregate_predictions
         self._last_window_inds_ = None
-        super().__init__(module,
-                         *args,
-                         callbacks=callbacks,
-                         iterator_train__shuffle=iterator_train__shuffle,
-                         iterator_train__drop_last=iterator_train__drop_last,
-                         **kwargs)
+        super().__init__(
+            module,
+            *args,
+            callbacks=callbacks,
+            iterator_train__shuffle=iterator_train__shuffle,
+            iterator_train__drop_last=iterator_train__drop_last,
+            **kwargs,
+        )
 
     def get_iterator(self, dataset, training=False, drop_index=True):
         iterator = super().get_iterator(dataset, training=training)
@@ -155,7 +164,9 @@ class EEGRegressor(_EEGNeuralNet, NeuralNetRegressor):
             warnings.warn(
                 "This method was designed to predict trials in cropped mode. "
                 "Calling it when cropped is False will give the same result as "
-                "'.predict'.", UserWarning)
+                "'.predict'.",
+                UserWarning,
+            )
             preds = self.predict(X)
             if return_targets:
                 return preds, np.concatenate([X[i][1] for i in range(len(X))])

braindecode/samplers/__init__.py

@@ -1,10 +1,18 @@
-"""Classes to sample examples.
-"""
+"""Classes to sample examples."""
 
-from .base import RecordingSampler, SequenceSampler, BalancedSequenceSampler
-from .ssl import RelativePositioningSampler
+from .base import (
+    BalancedSequenceSampler,
+    DistributedRecordingSampler,
+    RecordingSampler,
+    SequenceSampler,
+)
+from .ssl import DistributedRelativePositioningSampler, RelativePositioningSampler
 
-__all__ = ["RecordingSampler",
-           "SequenceSampler",
-           "BalancedSequenceSampler",
-           "RelativePositioningSampler"]
+__all__ = [
+    "RecordingSampler",
+    "SequenceSampler",
+    "BalancedSequenceSampler",
+    "RelativePositioningSampler",
+    "DistributedRecordingSampler",
+    "DistributedRelativePositioningSampler",
+]

braindecode/samplers/base.py

@@ -4,17 +4,120 @@ Sampler classes.
 
 # Authors: Hubert Banville <hubert.jbanville@gmail.com>
 #          Theo Gnassounou <>
+#          Young Truong <dt.young112@gmail.com>
 #
 # License: BSD (3-clause)
 
+from typing import Optional
+
 import numpy as np
-from torch.utils.data.sampler import Sampler
 from sklearn.utils import check_random_state
+from torch.utils.data.distributed import DistributedSampler
+from torch.utils.data.sampler import Sampler
 
 
 class RecordingSampler(Sampler):
     """Base sampler simplifying sampling from recordings.
 
+    Parameters
+    ----------
+    metadata : pd.DataFrame
+        DataFrame with at least one of {subject, session, run} columns for each
+        window in the BaseConcatDataset to sample examples from. Normally
+        obtained with `BaseConcatDataset.get_metadata()`. For instance,
+        `metadata.head()` might look like this:
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+        | i_window_in_trial | i_start_in_trial| i_stop_in_trial | target | subject  | session   | run   |
+        +===================+=================+=================+========+==========+===========+=======+
+        | 0                 | 0               | 500             | -1     | 4        | session_T | run_0 |
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+        | 1                 | 500             | 1000            | -1     | 4        | session_T | run_0 |
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+        | 2                 | 1000            | 1500            | -1     | 4        | session_T | run_0 |
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+        | 3                 | 1500            | 2000            | -1     | 4        | session_T | run_0 |
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+        | 4                 | 2000            | 2500            | -1     | 4        | session_T | run_0 |
+        +-------------------+-----------------+-----------------+--------+----------+-----------+-------+
+
+    random_state : np.RandomState | int | None
+        Random state.
+
+    Attributes
+    ----------
+    info : pd.DataFrame
+        Series with MultiIndex index which contains the subject, session, run
+        and window indices information in an easily accessible structure for
+        quick sampling of windows.
+    n_recordings : int
+        Number of recordings available.
+    """
+
+    def __init__(self, metadata, random_state=None):
+        self.metadata = metadata
+        self.info = self._init_info(metadata)
+        self.rng = check_random_state(random_state)
+
+    def _init_info(self, metadata, required_keys=None):
+        """Initialize ``info`` DataFrame.
+
+        Parameters
+        ----------
+        required_keys : list(str) | None
+            List of additional columns of the metadata DataFrame that we should
+            groupby when creating ``info``.
+
+        Returns
+        -------
+            See class attributes.
+        """
+        keys = [k for k in ["subject", "session", "run"] if k in self.metadata.columns]
+        if not keys:
+            raise ValueError(
+                "metadata must contain at least one of the following columns: "
+                "subject, session or run."
+            )
+
+        if required_keys is not None:
+            missing_keys = [k for k in required_keys if k not in self.metadata.columns]
+            if len(missing_keys) > 0:
+                raise ValueError(f"Columns {missing_keys} were not found in metadata.")
+            keys += required_keys
+
+        metadata = metadata.reset_index().rename(columns={"index": "window_index"})
+        info = (
+            metadata.reset_index()
+            .groupby(keys)[["index", "i_start_in_trial"]]
+            .agg(["unique"])
+        )
+        info.columns = info.columns.get_level_values(0)
+
+        return info
+
+    def sample_recording(self):
+        """Return a random recording index."""
+        # XXX docstring missing
+        return self.rng.choice(self.n_recordings)
+
+    def sample_window(self, rec_ind=None):
+        """Return a specific window."""
+        # XXX docstring missing
+        if rec_ind is None:
+            rec_ind = self.sample_recording()
+        win_ind = self.rng.choice(self.info.iloc[rec_ind]["index"])
+        return win_ind, rec_ind
+
+    def __iter__(self):
+        raise NotImplementedError
+
+    @property
+    def n_recordings(self):
+        return self.info.shape[0]
+
+
+class DistributedRecordingSampler(DistributedSampler):
+    """Base sampler simplifying sampling from recordings in distributed setting.
+
     Parameters
     ----------
     metadata : pd.DataFrame
@@ -41,11 +144,22 @@ class RecordingSampler(Sampler):
         quick sampling of windows.
     n_recordings : int
         Number of recordings available.
+    kwargs : dict
+        Additional keyword arguments to pass to torch DistributedSampler.
+        See https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
     """
-    def __init__(self, metadata, random_state=None):
+
+    def __init__(
+        self,
+        metadata,
+        random_state=None,
+        **kwargs,
+    ):
         self.metadata = metadata
         self.info = self._init_info(metadata)
         self.rng = check_random_state(random_state)
+        # send information to DistributedSampler parent to handle data splitting among workers
+        super().__init__(self.info, seed=random_state, **kwargs)
 
     def _init_info(self, metadata, required_keys=None):
         """Initialize ``info`` DataFrame.
@@ -60,50 +174,48 @@ class RecordingSampler(Sampler):
         -------
             See class attributes.
         """
-        keys = [k for k in ['subject', 'session', 'run']
-                if k in self.metadata.columns]
+        keys = [k for k in ["subject", "session", "run"] if k in self.metadata.columns]
         if not keys:
             raise ValueError(
-                'metadata must contain at least one of the following columns: '
-                'subject, session or run.')
+                "metadata must contain at least one of the following columns: "
+                "subject, session or run."
+            )
 
         if required_keys is not None:
-            missing_keys = [
-                k for k in required_keys if k not in self.metadata.columns]
+            missing_keys = [k for k in required_keys if k not in self.metadata.columns]
             if len(missing_keys) > 0:
-                raise ValueError(
-                    f'Columns {missing_keys} were not found in metadata.')
+                raise ValueError(f"Columns {missing_keys} were not found in metadata.")
             keys += required_keys
 
-        metadata = metadata.reset_index().rename(
-            columns={'index': 'window_index'})
-        info = metadata.reset_index().groupby(keys)[
-            ['index', 'i_start_in_trial']].agg(['unique'])
+        metadata = metadata.reset_index().rename(columns={"index": "window_index"})
+        info = (
+            metadata.reset_index()
+            .groupby(keys)[["index", "i_start_in_trial"]]
+            .agg(["unique"])
+        )
         info.columns = info.columns.get_level_values(0)
 
         return info
 
     def sample_recording(self):
         """Return a random recording index.
+        super().__iter__() contains indices of datasets specific to the current process
+        determined by the DistributedSampler
         """
         # XXX docstring missing
-        return self.rng.choice(self.n_recordings)
+        return self.rng.choice(list(super().__iter__()))
 
     def sample_window(self, rec_ind=None):
-        """Return a specific window.
-        """
+        """Return a specific window."""
         # XXX docstring missing
         if rec_ind is None:
             rec_ind = self.sample_recording()
-        win_ind = self.rng.choice(self.info.iloc[rec_ind]['index'])
+        win_ind = self.rng.choice(self.info.iloc[rec_ind]["index"])
         return win_ind, rec_ind
 
-    def __iter__(self):
-        raise NotImplementedError
-
     @property
     def n_recordings(self):
-        return self.info.shape[0]
+        return super().__len__()
 
 
 class SequenceSampler(RecordingSampler):
@@ -131,8 +243,10 @@ class SequenceSampler(RecordingSampler):
         Array of shape (n_sequences,) that indicates from which file each
         sequence comes from. Useful e.g. to do self-ensembling.
     """
-    def __init__(self, metadata, n_windows, n_windows_stride, randomize=False,
-                 random_state=None):
+
+    def __init__(
+        self, metadata, n_windows, n_windows_stride, randomize=False, random_state=None
+    ):
         super().__init__(metadata, random_state=random_state)
         self.randomize = randomize
         self.n_windows = n_windows
@@ -152,8 +266,11 @@ class SequenceSampler(RecordingSampler):
             each sequence. Useful e.g. to do self-ensembling.
         """
         end_offset = 1 - self.n_windows if self.n_windows > 1 else None
-        start_inds = self.info['index'].apply(
-            lambda x: x[:end_offset:self.n_windows_stride]).values
+        start_inds = (
+            self.info["index"]
+            .apply(lambda x: x[: end_offset : self.n_windows_stride])
+            .values
+        )
         file_ids = [[i] * len(inds) for i, inds in enumerate(start_inds)]
         return np.concatenate(start_inds), np.concatenate(file_ids)
 
@@ -200,12 +317,13 @@ class BalancedSequenceSampler(RecordingSampler):
            Med. 4, 72 (2021).
            https://github.com/perslev/U-Time/blob/master/utime/models/usleep.py
     """
+
     def __init__(self, metadata, n_windows, n_sequences=10, random_state=None):
         super().__init__(metadata, random_state=random_state)
 
         self.n_windows = n_windows
         self.n_sequences = n_sequences
-        self.info_class = self._init_info(metadata, required_keys=['target'])
+        self.info_class = self._init_info(metadata, required_keys=["target"])
 
     def sample_class(self, rec_ind=None):
         """Return a random class.
@@ -225,8 +343,7 @@ class BalancedSequenceSampler(RecordingSampler):
         """
         if rec_ind is None:
             rec_ind = self.sample_recording()
-        available_classes = self.info_class.loc[
-            self.info.iloc[rec_ind].name].index
+        available_classes = self.info_class.loc[self.info.iloc[rec_ind].name].index
         return self.rng.choice(available_classes), rec_ind
 
     def _sample_seq_start_ind(self, rec_ind=None, class_ind=None):
@@ -257,15 +374,14 @@ class BalancedSequenceSampler(RecordingSampler):
         if class_ind is None:
             class_ind, rec_ind = self.sample_class(rec_ind)
 
-        rec_inds = self.info.iloc[rec_ind]['index']
+        rec_inds = self.info.iloc[rec_ind]["index"]
         len_rec_inds = len(rec_inds)
 
         row = self.info.iloc[rec_ind].name
         if not isinstance(row, tuple):
             # Theres's only one category, e.g. "subject"
             row = tuple([row])
-        available_indices = self.info_class.loc[
-            row + tuple([class_ind]), 'index']
+        available_indices = self.info_class.loc[row + tuple([class_ind]), "index"]
         win_ind = self.rng.choice(available_indices)
         win_ind_in_rec = np.where(rec_inds == win_ind)[0][0]
 

braindecode/samplers/ssl.py

@@ -3,12 +3,16 @@ Self-supervised learning samplers.
 """
 
 # Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#          Young Truong <dt.young112@gmail.com>
 #
 # License: BSD (3-clause)
 
+import warnings
+
 import numpy as np
+import torch.distributed as dist
 
-from . import RecordingSampler
+from . import DistributedRecordingSampler, RecordingSampler
 
 
 class RelativePositioningSampler(RecordingSampler):
@@ -45,8 +49,17 @@ class RelativePositioningSampler(RecordingSampler):
            signals with self-supervised learning.
            arXiv preprint arXiv:2007.16104.
     """
-    def __init__(self, metadata, tau_pos, tau_neg, n_examples, tau_max=None,
-                 same_rec_neg=True, random_state=None):
+
+    def __init__(
+        self,
+        metadata,
+        tau_pos,
+        tau_neg,
+        n_examples,
+        tau_max=None,
+        same_rec_neg=True,
+        random_state=None,
+    ):
         super().__init__(metadata, random_state=random_state)
 
         self.tau_pos = tau_pos
@@ -56,25 +69,153 @@ class RelativePositioningSampler(RecordingSampler):
         self.same_rec_neg = same_rec_neg
 
         if not same_rec_neg and self.n_recordings < 2:
-            raise ValueError('More than one recording must be available when '
-                             'using across-recording negative sampling.')
+            raise ValueError(
+                "More than one recording must be available when "
+                "using across-recording negative sampling."
+            )
 
     def _sample_pair(self):
-        """Sample a pair of two windows.
+        """Sample a pair of two windows."""
+        # Sample first window
+        win_ind1, rec_ind1 = self.sample_window()
+        ts1 = self.metadata.iloc[win_ind1]["i_start_in_trial"]
+        ts = self.info.iloc[rec_ind1]["i_start_in_trial"]
+
+        # Decide whether the pair will be positive or negative
+        pair_type = self.rng.binomial(1, 0.5)
+        win_ind2 = None
+        if pair_type == 0:  # Negative example
+            if self.same_rec_neg:
+                mask = ((ts <= ts1 - self.tau_neg) & (ts >= ts1 - self.tau_max)) | (
+                    (ts >= ts1 + self.tau_neg) & (ts <= ts1 + self.tau_max)
+                )
+            else:
+                rec_ind2 = rec_ind1
+                while rec_ind2 == rec_ind1:
+                    win_ind2, rec_ind2 = self.sample_window()
+        elif pair_type == 1:  # Positive example
+            mask = (ts >= ts1 - self.tau_pos) & (ts <= ts1 + self.tau_pos)
+
+        if win_ind2 is None:
+            mask[ts == ts1] = False  # same window cannot be sampled twice
+            if sum(mask) == 0:
+                raise NotImplementedError
+            win_ind2 = self.rng.choice(self.info.iloc[rec_ind1]["index"][mask])
+
+        return win_ind1, win_ind2, float(pair_type)
+
+    def presample(self):
+        """Presample examples.
+
+        Once presampled, the examples are the same from one epoch to another.
+        """
+        self.examples = [self._sample_pair() for _ in range(self.n_examples)]
+        return self
+
+    def __iter__(self):
+        """
+        Iterate over pairs.
+
+        Yields
+        ------
+        int
+            Position of the first window in the dataset.
+        int
+            Position of the second window in the dataset.
+        float
+            0 for a negative pair, 1 for a positive pair.
         """
+        for i in range(self.n_examples):
+            if hasattr(self, "examples"):
+                yield self.examples[i]
+            else:
+                yield self._sample_pair()
+
+    def __len__(self):
+        return self.n_examples
+
+
+class DistributedRelativePositioningSampler(DistributedRecordingSampler):
+    """Sample examples for the relative positioning task from [Banville2020]_ in distributed mode.
+
+    Sample examples as tuples of two window indices, with a label indicating
+    whether the windows are close or far, as defined by tau_pos and tau_neg.
+
+    Parameters
+    ----------
+    metadata : pd.DataFrame
+        See RecordingSampler.
+    tau_pos : int
+        Size of the positive context, in samples. A positive pair contains two
+        windows x1 and x2 which are separated by at most `tau_pos` samples.
+    tau_neg : int
+        Size of the negative context, in samples. A negative pair contains two
+        windows x1 and x2 which are separated by at least `tau_neg` samples and
+        at most `tau_max` samples. Ignored if `same_rec_neg` is False.
+    n_examples : int
+        Number of pairs to extract.
+    tau_max : int | None
+        See `tau_neg`.
+    same_rec_neg : bool
+        If True, sample negative pairs from within the same recording. If
+        False, sample negative pairs from two different recordings.
+    random_state : None | np.RandomState | int
+        Random state.
+    kwargs: dict
+        Additional keyword arguments to pass to torch DistributedSampler.
+        See https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
+
+    References
+    ----------
+    .. [Banville2020] Banville, H., Chehab, O., Hyvärinen, A., Engemann, D. A.,
+        & Gramfort, A. (2020). Uncovering the structure of clinical EEG
+        signals with self-supervised learning.
+        arXiv preprint arXiv:2007.16104.
+    """
+
+    def __init__(
+        self,
+        metadata,
+        tau_pos,
+        tau_neg,
+        n_examples,
+        tau_max=None,
+        same_rec_neg=True,
+        random_state=None,
+        **kwargs,
+    ):
+        super().__init__(metadata, random_state=random_state, **kwargs)
+        self.tau_pos = tau_pos
+        self.tau_neg = tau_neg
+        self.tau_max = np.inf if tau_max is None else tau_max
+        self.same_rec_neg = same_rec_neg
+
+        self.n_examples = n_examples // self.info.shape[0] * self.n_recordings
+        warnings.warn(
+            f"Rank {dist.get_rank()} - Number of datasets: {self.n_recordings}"
+        )
+        warnings.warn(f"Rank {dist.get_rank()} - Number of samples: {self.n_examples}")
+
+        if not same_rec_neg and self.n_recordings < 2:
+            raise ValueError(
+                "More than one recording must be available when "
+                "using across-recording negative sampling."
+            )
+
+    def _sample_pair(self):
+        """Sample a pair of two windows."""
         # Sample first window
         win_ind1, rec_ind1 = self.sample_window()
-        ts1 = self.metadata.iloc[win_ind1]['i_start_in_trial']
-        ts = self.info.iloc[rec_ind1]['i_start_in_trial']
+        ts1 = self.metadata.iloc[win_ind1]["i_start_in_trial"]
+        ts = self.info.iloc[rec_ind1]["i_start_in_trial"]
 
         # Decide whether the pair will be positive or negative
         pair_type = self.rng.binomial(1, 0.5)
         win_ind2 = None
         if pair_type == 0:  # Negative example
             if self.same_rec_neg:
-                mask = (
-                    ((ts <= ts1 - self.tau_neg) & (ts >= ts1 - self.tau_max)) |
-                    ((ts >= ts1 + self.tau_neg) & (ts <= ts1 + self.tau_max))
+                mask = ((ts <= ts1 - self.tau_neg) & (ts >= ts1 - self.tau_max)) | (
+                    (ts >= ts1 + self.tau_neg) & (ts <= ts1 + self.tau_max)
                 )
             else:
                 rec_ind2 = rec_ind1
@@ -87,7 +228,7 @@ class RelativePositioningSampler(RecordingSampler):
             mask[ts == ts1] = False  # same window cannot be sampled twice
             if sum(mask) == 0:
                 raise NotImplementedError
-            win_ind2 = self.rng.choice(self.info.iloc[rec_ind1]['index'][mask])
+            win_ind2 = self.rng.choice(self.info.iloc[rec_ind1]["index"][mask])
 
         return win_ind1, win_ind2, float(pair_type)
 
@@ -100,16 +241,20 @@ class RelativePositioningSampler(RecordingSampler):
         return self
 
     def __iter__(self):
-        """Iterate over pairs.
+        """
+        Iterate over pairs.
 
         Yields
         ------
-            (int): position of the first window in the dataset.
-            (int): position of the second window in the dataset.
-            (float): 0 for negative pair, 1 for positive pair.
+        int
+            Position of the first window in the dataset.
+        int
+            Position of the second window in the dataset.
+        float
+            0 for a negative pair, 1 for a positive pair.
         """
         for i in range(self.n_examples):
-            if hasattr(self, 'examples'):
+            if hasattr(self, "examples"):
                 yield self.examples[i]
             else:
                 yield self._sample_pair()

braindecode/training/__init__.py

@@ -2,14 +2,22 @@
 Functionality for skorch-based training.
 """
 
+from .losses import CroppedLoss, TimeSeriesLoss, mixup_criterion
+from .scoring import (
+    CroppedTimeSeriesEpochScoring,
+    CroppedTrialEpochScoring,
+    PostEpochTrainScoring,
+    predict_trials,
+    trial_preds_from_window_preds,
+)
 
-from .losses import CroppedLoss, mixup_criterion, TimeSeriesLoss
-from .scoring import (CroppedTrialEpochScoring, PostEpochTrainScoring,
-                      CroppedTimeSeriesEpochScoring, trial_preds_from_window_preds, predict_trials)
-
-__all__ = ["CroppedLoss", "mixup_criterion", "TimeSeriesLoss",
-           "CroppedTrialEpochScoring",
-           "PostEpochTrainScoring",
-           "CroppedTimeSeriesEpochScoring",
-           "trial_preds_from_window_preds",
-           "predict_trials"]
+__all__ = [
+    "CroppedLoss",
+    "mixup_criterion",
+    "TimeSeriesLoss",
+    "CroppedTrialEpochScoring",
+    "PostEpochTrainScoring",
+    "CroppedTimeSeriesEpochScoring",
+    "trial_preds_from_window_preds",
+    "predict_trials",
+]

braindecode/training/callbacks.py

@@ -2,8 +2,8 @@
 #
 # License: BSD (3-clause)
 
-from skorch.callbacks import Callback
 import torch
+from skorch.callbacks import Callback
 
 
 class MaxNormConstraintCallback(Callback):
@@ -20,6 +20,4 @@ class MaxNormConstraintCallback(Callback):
                     )
                     last_weight = module.weight
             if last_weight is not None:
-                last_weight.data = torch.renorm(
-                    last_weight.data, 2, 0, maxnorm=0.5
-                )
+                last_weight.data = torch.renorm(last_weight.data, 2, 0, maxnorm=0.5)

braindecode/training/losses.py

@@ -96,15 +96,10 @@ def mixup_criterion(preds, target):
         # unpack target
         y_a, y_b, lam = target
         # compute loss per sample
-        loss_a = torch.nn.functional.nll_loss(preds,
-                                              y_a,
-                                              reduction='none')
-        loss_b = torch.nn.functional.nll_loss(preds,
-                                              y_b,
-                                              reduction='none')
+        loss_a = torch.nn.functional.nll_loss(preds, y_a, reduction="none")
+        loss_b = torch.nn.functional.nll_loss(preds, y_b, reduction="none")
         # compute weighted mean
         ret = torch.mul(lam, loss_a) + torch.mul(1 - lam, loss_b)
         return ret.mean()
     else:
-        return torch.nn.functional.nll_loss(preds,
-                                            target)
+        return torch.nn.functional.nll_loss(preds, target)