Add torch nn

Author: hoonyyhoon

train.py

@@ -11,6 +11,7 @@
 import numpy as np
 import torch.distributed as dist
 import torch.nn.functional as F
+import torch.nn as nn
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler
 import torch.utils.data
@@ -103,7 +104,6 @@ def train(hyp, opt, device, tb_writer=None, wandb=None):
     for k, v in model.named_parameters():
         v.requires_grad = True
 
-
     if opt.adam:
         optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
     else: