diff --git a/train.py b/train.py
index 89c5c9f..8700078 100644
--- a/train.py
+++ b/train.py
@@ -10,7 +10,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.optim as optim
 import sys
-import re
 
 class Net(nn.Module):
     def __init__(self):
@@ -76,10 +75,9 @@ if __name__ == '__main__':
     optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
 
     epochs = sys.argv[1]
+    print(epochs)
 
-    epochs_m = re.findall(r'\d+\.\d+', epochs)
-
-    trainNet(trainloader, criterion, optimizer, int(float(epochs_m[0])))
+    trainNet(trainloader, criterion, optimizer, int(float(epochs)))
 
     PATH = './cifar_net.pth'
     torch.save(net.state_dict(), PATH)