Apply the changes (#3069)

ignite/engine/__init__.py

@@ -44,7 +44,7 @@ def _prepare_batch(
 def supervised_training_step(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
-    loss_fn: Union[Callable, torch.nn.Module],
+    loss_fn: Union[Callable[[Any, Any], torch.Tensor], torch.nn.Module],
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
@@ -57,7 +57,7 @@ def supervised_training_step(
     Args:
         model: the model to train.
         optimizer: the optimizer to use.
-        loss_fn: the loss function to use.
+        loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
         device: device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
             Device can be CPU, GPU.
@@ -120,7 +120,7 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
 def supervised_training_step_amp(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
-    loss_fn: Union[Callable, torch.nn.Module],
+    loss_fn: Union[Callable[[Any, Any], torch.Tensor], torch.nn.Module],
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
@@ -134,7 +134,7 @@ def supervised_training_step_amp(
     Args:
         model: the model to train.
         optimizer: the optimizer to use.
-        loss_fn: the loss function to use.
+        loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
         device: device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
             Device can be CPU, GPU.
@@ -212,7 +212,7 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
 def supervised_training_step_apex(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
-    loss_fn: Union[Callable, torch.nn.Module],
+    loss_fn: Union[Callable[[Any, Any], torch.Tensor], torch.nn.Module],
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
@@ -225,7 +225,7 @@ def supervised_training_step_apex(
     Args:
         model: the model to train.
         optimizer: the optimizer to use.
-        loss_fn: the loss function to use.
+        loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
         device: device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
             Device can be CPU, GPU.
@@ -295,7 +295,7 @@ def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[to
 def supervised_training_step_tpu(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
-    loss_fn: Union[Callable, torch.nn.Module],
+    loss_fn: Union[Callable[[Any, Any], torch.Tensor], torch.nn.Module],
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
@@ -308,7 +308,7 @@ def supervised_training_step_tpu(
     Args:
         model: the model to train.
         optimizer: the optimizer to use.
-        loss_fn: the loss function to use.
+        loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
         device: device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
             Device can be CPU, TPU.
@@ -404,7 +404,7 @@ def _check_arg(
 def create_supervised_trainer(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
-    loss_fn: Union[Callable, torch.nn.Module],
+    loss_fn: Union[Callable[[Any, Any], torch.Tensor], torch.nn.Module],
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
@@ -420,7 +420,7 @@ def create_supervised_trainer(
     Args:
         model: the model to train.
         optimizer: the optimizer to use.
-        loss_fn: the loss function to use.
+        loss_fn: the loss function that receives `y_pred` and `y`, and returns the loss as a tensor.
         device: device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
             Device can be CPU, GPU or TPU.

tests/ignite/conftest.py

@@ -397,6 +397,7 @@ def gloo_hvd_executor():
             ],
         ),
     ],
+    scope="class",
 )
 def distributed(request, local_rank, world_size):
     if request.param in ("nccl", "gloo_cpu", "gloo"):