Imported from https://github.com/bbclr20/PyTorch-Examples