Examples

TrainUnit Example

The TrainUnit example shows how to use the TrainUnit to train a basic model.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/train_unit_example.py

AutoUnit Example

The AutoUnit example shows how to use the AutoUnit to train a basic model with less code, and more training features enabled out of the box.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/auto_unit_example.py

TorchData Train Example

The TorchData Train example shows how to use TorchData’s DataPipe and Dataloader2. when training a basic model with TNT.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/torchdata_train_example.py

MNIST Example

The MNIST example shows how to use TNT to train a convnet model on the MNIST dataset.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/mnist/main.py

MinGPT Example

The MinGPT example shows how to use the AutoUnit to train a MinGPT model.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/mingpt/main.py

TorchRec Example

The TorchRec example shows how to use TrainUnit with TorchRec.

It can be found here: https://github.com/pytorch/tnt/blob/master/examples/torchrec/main.py