def wrap_model(self, model: Module) -> Module:
        """
        Wraps a given model to enable distributed training.

        The given model will be wrapped using :class:`DistributedDataParallel`.

        :return: The model wrapped in :class:`DistributedDataParallel` if
            running a distributed training, or the model itself if running a
            single-process training.
        """
        # Note: find_unused_parameters is needed for multi task models.
        if self.is_distributed:
            if self.forced_cuda_comm or self.use_cuda:
                # forced_cuda_comm is True if using NCCL; use_cuda may be true
                # even when not using NCCL.
                # User already warned if using NCCL with use_cuda==False.
                # device_ids must be a single device id
                # (an int, a device object or a str)
                # If not set, output_device defaults to device_ids[0]
                return DistributedDataParallel(
                    model, device_ids=[self.make_device()], find_unused_parameters=True
                )
            else:
                return DistributedDataParallel(model, find_unused_parameters=True)
        else:
            return model

        Wraps a given model to enable distributed training.

        The given model will be wrapped using :class:`DistributedDataParallel`.

        :return: The model wrapped in :class:`DistributedDataParallel` if
            running a distributed training, or the model itself if running a
            single-process training.
        
avalanche/distributed/distributed_helper.py