Minor import and spelling updates (#1133)

train.py

@@ -1,12 +1,12 @@
 import argparse
 import logging
-import math
 import os
 import random
 import shutil
 import time
 from pathlib import Path
 
+import math
 import numpy as np
 import torch.distributed as dist
 import torch.nn.functional as F