Repository: pytorch/ignite

Issue:
Add more options to create_supervised_trainer
## 🚀 Feature

Idea is to add more options to `create_supervised_trainer` helper:
- native AMP option similar to https://github.com/lidq92/LinearityIQA/blob/master/modified_ignite_engine.py 
  - we may either use nvidia/apex if installed or torch native amp if v1.6.0. Otherwise, raise an error.
- grads cliping as proposed #419


For Hacktoberfest/PyDataGlobal contributors, feel free to ask questions for details if any and say that you would like to tackle the issue.
Please, take a look at [CONTRIBUTING guide](https://github.com/pytorch/ignite/blob/master/CONTRIBUTING.md).


Discussion:
I would like to work on this. 
@apthagowda97 sure, please go ahead. I assign you to this issue
@apthagowda97 any updates on this ? Anyway, do not hesitate to ask question if something is not clear about the issue.
By going through issue #419 its best to add gradient clipping to `contrib.engines` as discussed compared with changing `create_supervised_trainer`. Should I name it `create_supervised_gc_trainer`?
@apthagowda97 I think it is better let's keep the same name and add the options. Such that there is one multi-option helper function.
@apthagowda97 any updates on this issue ?
Sorry I was busy from past few days. I ll try to finish it Asap. 

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -10,10 +10,6 @@
 from ignite.metrics import Metric
 from ignite.utils import convert_tensor
 
-if idist.has_xla_support:
-    import torch_xla.core.xla_model as xm
-
-
 __all__ = [
     "State",
     "create_supervised_trainer",
@@ -25,6 +21,10 @@
     "EventEnum",
     "CallableEventWithFilter",
     "RemovableEventHandle",
+    "supervised_training_step",
+    "supervised_training_step_amp",
+    "supervised_training_step_apex",
+    "supervised_training_step_tpu",
 ]
 
 
@@ -41,6 +41,233 @@ def _prepare_batch(
     )
 
 
+def supervised_training_step(
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: Union[Callable, torch.nn.Module],
+    device: Optional[Union[str, torch.device]] = None,
+    non_blocking: bool = False,
+    prepare_batch: Callable = _prepare_batch,
+    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+) -> Callable:
+    """Factory function for supervised training.
+
+    Args:
+        model (torch.nn.Module): the model to train.
+        optimizer (torch.optim.Optimizer): the optimizer to use.
+        loss_fn (torch.nn loss function): the loss function to use.
+        device (str, optional): device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+            Device can be CPU, GPU.
+        non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            with respect to the host. For other cases, this argument has no effect.
+        prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
+            tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): function that receives 'x', 'y', 'y_pred', 'loss' and returns value
+            to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
+
+    Returns:
+        Callable: update function.
+
+    .. versionadded:: 0.5.0
+    """
+
+    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        model.train()
+        optimizer.zero_grad()
+        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
+        y_pred = model(x)
+        loss = loss_fn(y_pred, y)
+        loss.backward()
+        optimizer.step()
+        return output_transform(x, y, y_pred, loss)
+
+    return update
+
+
+def supervised_training_step_amp(
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: Union[Callable, torch.nn.Module],
+    device: Optional[Union[str, torch.device]] = None,
+    non_blocking: bool = False,
+    prepare_batch: Callable = _prepare_batch,
+    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+    scaler: Optional["torch.cuda.amp.GradScaler"] = None,
+) -> Callable:
+    """Factory function for supervised training using ``torch.cuda.amp``.
+
+    Args:
+        model (torch.nn.Module): the model to train.
+        optimizer (torch.optim.Optimizer): the optimizer to use.
+        loss_fn (torch.nn loss function): the loss function to use.
+        device (str, optional): device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+            Device can be CPU, GPU.
+        non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            with respect to the host. For other cases, this argument has no effect.
+        prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
+            tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): function that receives 'x', 'y', 'y_pred', 'loss' and returns value
+            to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
+        scaler (torch.cuda.amp.GradScaler, optional): GradScaler instance for gradient scaling. (default: None)
+
+    Returns:
+        Callable: update function
+
+    .. versionadded:: 0.5.0
+    """
+
+    try:
+        from torch.cuda.amp import autocast
+    except ImportError:
+        raise ImportError("Please install torch>=1.6.0 to use amp_mode='amp'.")
+
+    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        model.train()
+        optimizer.zero_grad()
+        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
+        with autocast(enabled=True):
+            y_pred = model(x)
+            loss = loss_fn(y_pred, y)
+        if scaler:
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            loss.backward()
+            optimizer.step()
+        return output_transform(x, y, y_pred, loss)
+
+    return update
+
+
+def supervised_training_step_apex(
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: Union[Callable, torch.nn.Module],
+    device: Optional[Union[str, torch.device]] = None,
+    non_blocking: bool = False,
+    prepare_batch: Callable = _prepare_batch,
+    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+) -> Callable:
+    """Factory function for supervised training using apex.
+
+    Args:
+        model (torch.nn.Module): the model to train.
+        optimizer (torch.optim.Optimizer): the optimizer to use.
+        loss_fn (torch.nn loss function): the loss function to use.
+        device (str, optional): device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+            Device can be CPU, GPU.
+        non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            with respect to the host. For other cases, this argument has no effect.
+        prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
+            tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): function that receives 'x', 'y', 'y_pred', 'loss' and returns value
+            to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
+
+    Returns:
+        Callable: update function.
+
+    .. versionadded:: 0.5.0
+    """
+
+    try:
+        from apex import amp as apex_amp
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("Please install apex from https://github.com/nvidia/apex to use amp_mode='apex'.")
+
+    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        model.train()
+        optimizer.zero_grad()
+        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
+        y_pred = model(x)
+        loss = loss_fn(y_pred, y)
+        with apex_amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+        optimizer.step()
+        return output_transform(x, y, y_pred, loss)
+
+    return update
+
+
+def supervised_training_step_tpu(
+    model: torch.nn.Module,
+    optimizer: torch.optim.Optimizer,
+    loss_fn: Union[Callable, torch.nn.Module],
+    device: Optional[Union[str, torch.device]] = None,
+    non_blocking: bool = False,
+    prepare_batch: Callable = _prepare_batch,
+    output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
+) -> Callable:
+    """Factory function for supervised training using ``torch_xla``.
+
+    Args:
+        model (torch.nn.Module): the model to train.
+        optimizer (torch.optim.Optimizer): the optimizer to use.
+        loss_fn (torch.nn loss function): the loss function to use.
+        device (str, optional): device type specification (default: None).
+            Applies to batches after starting the engine. Model *will not* be moved.
+            Device can be CPU, TPU.
+        non_blocking (bool, optional): if True and this copy is between CPU and GPU, the copy may occur asynchronously
+            with respect to the host. For other cases, this argument has no effect.
+        prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
+            tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): function that receives 'x', 'y', 'y_pred', 'loss' and returns value
+            to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
+
+    Returns:
+        Callable: update function.
+
+    .. versionadded:: 0.5.0
+    """
+    try:
+        import torch_xla.core.xla_model as xm
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("torch_xla cannot be imported, please install PyTorch XLA.")
+
+    def update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
+        model.train()
+        optimizer.zero_grad()
+        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
+        y_pred = model(x)
+        loss = loss_fn(y_pred, y)
+        loss.backward()
+        xm.optimizer_step(optimizer, barrier=True)
+        return output_transform(x, y, y_pred, loss)
+
+    return update
+
+
+def _check_arg(
+    on_tpu: bool, amp_mode: Optional[str], scaler: Optional[Union[bool, "torch.cuda.amp.GradScaler"]]
+) -> Tuple[Optional[str], Optional["torch.cuda.amp.GradScaler"]]:
+    """Checking tpu, amp and GradScaler instance combinations."""
+    if on_tpu and not idist.has_xla_support:
+        raise RuntimeError("In order to run on TPU, please install PyTorch XLA")
+
+    if amp_mode and on_tpu:
+        raise ValueError("amp_mode cannot be used with xla device. Consider using amp_mode=None or device='cuda'.")
+
+    if scaler:
+        if amp_mode != "amp":
+            raise ValueError(f"scaler argument is {scaler}, but amp_mode is {amp_mode}. Consider using amp_mode='amp'.")
+        elif amp_mode == "amp" and isinstance(scaler, bool):
+            try:
+                from torch.cuda.amp import GradScaler
+            except ImportError:
+                raise ImportError("Please install torch>=1.6.0 to use scaler argument.")
+            scaler = GradScaler(enabled=True)
+
+    if on_tpu:
+        return "tpu", None
+    elif scaler and amp_mode == "amp":
+        return amp_mode, scaler  # type: ignore[return-value]
+    else:
+        return amp_mode, None
+
+
 def create_supervised_trainer(
     model: torch.nn.Module,
     optimizer: torch.optim.Optimizer,
@@ -50,12 +277,14 @@ def create_supervised_trainer(
     prepare_batch: Callable = _prepare_batch,
     output_transform: Callable = lambda x, y, y_pred, loss: loss.item(),
     deterministic: bool = False,
+    amp_mode: Optional[str] = None,
+    scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
 ) -> Engine:
     """Factory function for creating a trainer for supervised models.
 
     Args:
-        model (`torch.nn.Module`): the model to train.
-        optimizer (`torch.optim.Optimizer`): the optimizer to use.
+        model (torch.nn.Module): the model to train.
+        optimizer (torch.optim.Optimizer): the optimizer to use.
         loss_fn (torch.nn loss function): the loss function to use.
         device (str, optional): device type specification (default: None).
             Applies to batches after starting the engine. Model *will not* be moved.
@@ -69,48 +298,71 @@ def create_supervised_trainer(
         deterministic (bool, optional): if True, returns deterministic engine of type
             :class:`~ignite.engine.deterministic.DeterministicEngine`, otherwise :class:`~ignite.engine.engine.Engine`
             (default: False).
+        amp_mode (str, optional): can be ``amp`` or ``apex``, model and optimizer will be casted to float16 using
+            `torch.cuda.amp <https://pytorch.org/docs/stable/amp.html>`_ for ``amp`` and
+            using `apex <https://nvidia.github.io/apex>`_ for ``apex``. (default: None)
+        scaler (torch.cuda.amp.GradScaler, bool, optional): GradScaler instance for gradient scaling if `torch>=1.6.0`
+            and ``amp_mode`` is ``amp``. If ``amp_mode`` is ``apex``, this argument will be ignored.
+            If True, will create default GradScaler. If GradScaler instance is passed, it will be used instead.
+            (default: False)
+
+    Note:
+        If ``scaler`` is True, GradScaler instance will be created internally and trainer state has attribute named
+        ``scaler`` for that instance and can be used for saving and loading.
 
     Note:
         `engine.state.output` for this engine is defined by `output_transform` parameter and is the loss
         of the processed batch by default.
 
     .. warning::
-
         The internal use of `device` has changed.
         `device` will now *only* be used to move the input data to the correct device.
         The `model` should be moved by the user before creating an optimizer.
         For more information see:
 
         - `PyTorch Documentation <https://pytorch.org/docs/stable/optim.html#constructing-it>`_
-
         - `PyTorch's Explanation <https://github.com/pytorch/pytorch/issues/7844#issuecomment-503713840>`_
 
+    .. warning::
+        If ``amp_mode='apex'`` , the model(s) and optimizer(s) must be initialized beforehand
+        since ``amp.initialize`` should be called after you have finished constructing your model(s)
+        and optimizer(s), but before you send your model through any DistributedDataParallel wrapper.
+
+        See more: https://nvidia.github.io/apex/amp.html#module-apex.amp
+
     Returns:
         Engine: a trainer engine with supervised update function.
+
+    .. versionchanged:: 0.5.0
+
+        - Added ``amp_mode`` argument for automatic mixed precision.
+        - Added ``scaler`` argument for gradient scaling.
     """
 
     device_type = device.type if isinstance(device, torch.device) else device
     on_tpu = "xla" in device_type if device_type is not None else False
-
-    if on_tpu and not idist.has_xla_support:
-        raise RuntimeError("In order to run on TPU, please install PyTorch XLA")
-
-    def _update(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
-        model.train()
-        optimizer.zero_grad()
-        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-        y_pred = model(x)
-        loss = loss_fn(y_pred, y)
-        loss.backward()
-
-        if on_tpu:
-            xm.optimizer_step(optimizer, barrier=True)
-        else:
-            optimizer.step()
-
-        return output_transform(x, y, y_pred, loss)
+    mode, _scaler = _check_arg(on_tpu, amp_mode, scaler)
+
+    if mode == "amp":
+        _update = supervised_training_step_amp(
+            model, optimizer, loss_fn, device, non_blocking, prepare_batch, output_transform, _scaler
+        )
+    elif mode == "apex":
+        _update = supervised_training_step_apex(
+            model, optimizer, loss_fn, device, non_blocking, prepare_batch, output_transform
+        )
+    elif mode == "tpu":
+        _update = supervised_training_step_tpu(
+            model, optimizer, loss_fn, device, non_blocking, prepare_batch, output_transform
+        )
+    else:
+        _update = supervised_training_step(
+            model, optimizer, loss_fn, device, non_blocking, prepare_batch, output_transform
+        )
 
     trainer = Engine(_update) if not deterministic else DeterministicEngine(_update)
+    if _scaler and scaler and isinstance(scaler, bool):
+        trainer.state.scaler = _scaler  # type: ignore[attr-defined]
 
     return trainer
 


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -1,6 +1,8 @@
 import os
 from distutils.version import LooseVersion
-from typing import Optional
+from importlib.util import find_spec
+from typing import Optional, Union
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -10,12 +12,16 @@
 from torch.optim import SGD
 
 import ignite.distributed as idist
-from ignite.engine import create_supervised_evaluator, create_supervised_trainer
+from ignite.engine import create_supervised_evaluator, create_supervised_trainer, supervised_training_step_tpu
 from ignite.metrics import MeanSquaredError
 
 
 def _test_create_supervised_trainer(
-    model_device: Optional[str] = None, trainer_device: Optional[str] = None, trace: bool = False
+    model_device: Optional[str] = None,
+    trainer_device: Optional[str] = None,
+    trace: bool = False,
+    amp_mode: str = None,
+    scaler: Union[bool, "torch.cuda.amp.GradScaler"] = False,
 ):
     model = Linear(1, 1)
 
@@ -30,10 +36,23 @@ def _test_create_supervised_trainer(
         example_input = torch.randn(1, 1)
         model = torch.jit.trace(model, example_input)
 
-    trainer = create_supervised_trainer(model, optimizer, mse_loss, device=trainer_device)
+    if amp_mode == "apex" and model_device == trainer_device == "cuda":
+        from apex import amp
 
-    x = torch.tensor([[1.0], [2.0]])
-    y = torch.tensor([[3.0], [5.0]])
+        model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
+
+    trainer = create_supervised_trainer(
+        model,
+        optimizer,
+        mse_loss,
+        device=trainer_device,
+        output_transform=lambda x, y, y_pred, loss: (y_pred, loss.item()),
+        amp_mode=amp_mode,
+        scaler=scaler,
+    )
+
+    x = torch.tensor([[0.1], [0.2]])
+    y = torch.tensor([[0.3], [0.5]])
     data = [(x, y)]
 
     assert model.weight.data[0, 0].item() == approx(0.0)
@@ -42,9 +61,16 @@ def _test_create_supervised_trainer(
     if model_device == trainer_device or ((model_device == "cpu") ^ (trainer_device == "cpu")):
         state = trainer.run(data)
 
-        assert state.output == approx(17.0)
-        assert model.weight.data[0, 0].item() == approx(1.3)
-        assert model.bias.item() == approx(0.8)
+        assert state.output[-1] == approx(0.17), state.output[-1]
+        assert round(model.weight.data[0, 0].item(), 3) == approx(0.013), model.weight.item()
+        assert round(model.bias.item(), 3) == approx(0.08), model.bias.item()
+
+        if amp_mode == "amp":
+            assert state.output[0].dtype is torch.half
+            if scaler and isinstance(scaler, bool):
+                assert hasattr(state, "scaler")
+            else:
+                assert not hasattr(state, "scaler")
     else:
         if LooseVersion(torch.__version__) >= LooseVersion("1.7.0"):
             # This is broken in 1.6.0 but will be probably fixed with 1.7.0
@@ -105,12 +131,82 @@ def test_create_supervised_trainer_traced_with_cpu():
     _test_create_supervised_trainer(trainer_device="cpu", trace=True)
 
 
+@pytest.mark.skipif(find_spec("apex"), reason="Skip if APEX")
+def test_create_supervised_trainer_apex_error():
+    with pytest.raises(
+        ModuleNotFoundError, match="Please install apex from https://github.com/nvidia/apex to use amp_mode='apex'."
+    ):
+        _test_create_supervised_trainer(amp_mode="apex")
+
+
+@pytest.fixture
+def mock_torch_cuda_amp_module():
+    with patch.dict(
+        "sys.modules",
+        {"torch.cuda.amp": None, "torch.cuda.amp.grad_scaler": None, "torch.cuda.amp.autocast_mode": None},
+    ):
+        yield torch
+
+
+def test_create_supervised_trainer_amp_error(mock_torch_cuda_amp_module):
+    with pytest.raises(ImportError, match="Please install torch>=1.6.0 to use amp_mode='amp'."):
+        _test_create_supervised_trainer(amp_mode="amp")
+    with pytest.raises(ImportError, match="Please install torch>=1.6.0 to use scaler argument."):
+        _test_create_supervised_trainer(amp_mode="amp", scaler=True)
+
+
+@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion("1.6.0"), reason="Skip if < 1.6.0")
+def test_create_supervised_trainer_scaler_not_amp():
+    scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
+    with pytest.raises(ValueError, match=f"scaler argument is {scaler}, but amp_mode is None."):
+        _test_create_supervised_trainer(amp_mode=None, scaler=scaler)
+    with pytest.raises(ValueError, match="scaler argument is True, but amp_mode is None."):
+        _test_create_supervised_trainer(amp_mode=None, scaler=True)
+    with pytest.raises(ValueError, match="scaler argument is True, but amp_mode is apex."):
+        _test_create_supervised_trainer(amp_mode="apex", scaler=True)
+    with pytest.raises(ValueError, match=f"scaler argument is {scaler}, but amp_mode is apex."):
+        _test_create_supervised_trainer(amp_mode="apex", scaler=scaler)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
 def test_create_supervised_trainer_on_cuda():
     model_device = trainer_device = "cuda"
     _test_create_supervised_trainer(model_device=model_device, trainer_device=trainer_device)
 
 
+@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion("1.6.0"), reason="Skip if < 1.6.0")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_create_supervised_trainer_on_cuda_amp():
+    model_device = trainer_device = "cuda"
+    _test_create_supervised_trainer(model_device=model_device, trainer_device=trainer_device, amp_mode="amp")
+
+
+@pytest.mark.skipif(LooseVersion(torch.__version__) < LooseVersion("1.6.0"), reason="Skip if < 1.6.0")
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+def test_create_supervised_trainer_on_cuda_amp_scaler():
+    model_device = trainer_device = "cuda"
+    _test_create_supervised_trainer(
+        model_device=model_device, trainer_device=trainer_device, amp_mode="amp", scaler=True
+    )
+    scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
+    _test_create_supervised_trainer(
+        model_device=model_device, trainer_device=trainer_device, amp_mode="amp", scaler=scaler
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
+@pytest.mark.skipif(not find_spec("apex"), reason="Skip if no APEX")
+def test_create_supervised_trainer_on_cuda_apex():
+    model_device = trainer_device = "cuda"
+    _test_create_supervised_trainer(model_device=model_device, trainer_device=trainer_device, amp_mode="apex")
+
+
+@pytest.mark.skipif(idist.has_xla_support, reason="Skip if has PyTorch XLA package")
+def test_supervised_training_step_tpu_no_xla():
+    with pytest.raises(ModuleNotFoundError, match="torch_xla cannot be imported, please install PyTorch XLA."):
+        supervised_training_step_tpu(model=None, optimizer=None, loss_fn=None)
+
+
 @pytest.mark.skipif(idist.has_xla_support, reason="Skip if has PyTorch XLA package")
 def test_create_supervised_trainer_on_tpu_no_xla():
     model_device = "cpu"
@@ -127,6 +223,14 @@ def test_create_supervised_trainer_on_tpu():
     _test_create_supervised_trainer(model_device=model_device, trainer_device=trainer_device)
 
 
+@pytest.mark.tpu
+@pytest.mark.skipif(not idist.has_xla_support, reason="Skip if no PyTorch XLA package")
+def test_create_supervised_trainer_on_tpu_amp():
+    model_device = trainer_device = "xla"
+    with pytest.raises(ValueError, match="amp_mode cannot be used with xla device."):
+        _test_create_supervised_trainer(model_device=model_device, trainer_device=trainer_device, amp_mode="amp")
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Skip if no GPU")
 def test_create_supervised_trainer_on_cuda_with_model_on_cpu():
     _test_create_supervised_trainer(trainer_device="cuda")
diff --git a/tests/run_tpu_tests.sh b/tests/run_tpu_tests.sh
--- a/tests/run_tpu_tests.sh
+++ b/tests/run_tpu_tests.sh
@@ -6,5 +6,5 @@ pytest --cov ignite --cov-report term-missing --cov-report xml tests/ -vvv -m tp
 
 if [ -z ${NUM_TPU_WORKERS+x} ]; then
     export NUM_TPU_WORKERS=1
-    pytest --cov ignite --cov-report xml tests/ -vvv -m tpu
+    pytest --cov ignite --cov-append --cov-report term-missing --cov-report xml tests/ -vvv -m tpu
 fi