diff --git a/.gitignore b/.gitignore
index c60a390..3e00abc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -66,3 +66,6 @@ target/
 
 #Idea
 .idea*
+/data/
+/test/
+/.mypy_cache/
diff --git a/train.py b/train.py
index f3636bd..fb15a9f 100644
--- a/train.py
+++ b/train.py
@@ -2,19 +2,20 @@
 import os
 import sys
 import time
+
+import tabulate
 import torch
 import torch.nn.functional as F
 import torchvision
+
 import models
 import utils
-import tabulate
-
 
 parser = argparse.ArgumentParser(description='SGD/SWA training')
 parser.add_argument('--dir', type=str, default=None, required=True, help='training directory (default: None)')
 
 parser.add_argument('--dataset', type=str, default='CIFAR10', help='dataset name (default: CIFAR10)')
-parser.add_argument('--data_path', type=str, default=None, required=True, metavar='PATH',
+parser.add_argument('--data_path', type=str, default="./data", required=False, metavar='PATH',
                     help='path to datasets location (default: None)')
 parser.add_argument('--batch_size', type=int, default=128, metavar='N', help='input batch size (default: 128)')
 parser.add_argument('--num_workers', type=int, default=4, metavar='N', help='number of workers (default: 4)')
@@ -75,13 +76,12 @@
         pin_memory=True
     )
 }
-num_classes = max(train_set.train_labels) + 1
+num_classes = max(train_set.targets) + 1
 
 print('Preparing model')
 model = model_cfg.base(*model_cfg.args, num_classes=num_classes, **model_cfg.kwargs)
 model.cuda()
 
-
 if args.swa:
     print('SWA training')
     swa_model = model_cfg.base(*model_cfg.args, num_classes=num_classes, **model_cfg.kwargs)
diff --git a/utils.py b/utils.py
index d355164..430ee19 100644
--- a/utils.py
+++ b/utils.py
@@ -1,5 +1,7 @@
 import os
+
 import torch
+import tqdm
 
 
 def adjust_learning_rate(optimizer, lr):
@@ -23,9 +25,9 @@ def train_epoch(loader, model, criterion, optimizer):
 
     model.train()
 
-    for i, (input, target) in enumerate(loader):
-        input = input.cuda(async=True)
-        target = target.cuda(async=True)
+    for i, (input, target) in enumerate(tqdm.tqdm(loader)):
+        input = input.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
         input_var = torch.autograd.Variable(input)
         target_var = torch.autograd.Variable(target)
 
@@ -36,7 +38,7 @@ def train_epoch(loader, model, criterion, optimizer):
         loss.backward()
         optimizer.step()
 
-        loss_sum += loss.data[0] * input.size(0)
+        loss_sum += loss.item() * input.size(0)
         pred = output.data.max(1, keepdim=True)[1]
         correct += pred.eq(target_var.data.view_as(pred)).sum().item()
 
@@ -53,15 +55,15 @@ def eval(loader, model, criterion):
     model.eval()
 
     for i, (input, target) in enumerate(loader):
-        input = input.cuda(async=True)
-        target = target.cuda(async=True)
+        input = input.cuda(non_blocking=True)
+        target = target.cuda(non_blocking=True)
         input_var = torch.autograd.Variable(input)
         target_var = torch.autograd.Variable(target)
 
         output = model(input_var)
         loss = criterion(output, target_var)
 
-        loss_sum += loss.data[0] * input.size(0)
+        loss_sum += loss.item() * input.size(0)
         pred = output.data.max(1, keepdim=True)[1]
         correct += pred.eq(target_var.data.view_as(pred)).sum().item()
 
@@ -121,7 +123,7 @@ def bn_update(loader, model):
     model.apply(lambda module: _get_momenta(module, momenta))
     n = 0
     for input, _ in loader:
-        input = input.cuda(async=True)
+        input = input.cuda(non_blocking=True)
         input_var = torch.autograd.Variable(input)
         b = input_var.data.size(0)