Solved by vfdev-5 in post #2 
              
              
                Hi @odats 
You can create a trainer with a custom train_step (or update_fn) as here: 
https://pytorch.org/ignite/concepts.html#engine 
HTH