diff --git a/train.py b/train.py
index f9ef656..73362d6 100644
--- a/train.py
+++ b/train.py
@@ -72,7 +72,7 @@ def collatefn(data):
     writer = SummaryWriter(args.logdir, flush_secs=.5)
     device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
     
-    model = fasterrcnn_mobilenet_v3_large_fpn(num_classes=9)
+    model = fasterrcnn_mobilenet_v3_large_fpn(num_classes=10)
     model.to(device)
     optim = Adam(model.parameters(), lr=args.lr)
 
@@ -128,4 +128,4 @@ def collatefn(data):
         writer.add_scalar('val/mAP@50', mout['map_50'], it)
         writer.add_scalar('val/mAP@75', mout['map_75'], it)
         writer.add_figure('val/label', tset.get_figure(ins[0], tgs[0]['boxes'], tgs[0]['labels']), it, close=True)
-        writer.add_figure('val/pred', tset.get_figure(ins[0], out[0]['boxes'], out[0]['labels']), it, close=True)
\ No newline at end of file
+        writer.add_figure('val/pred', tset.get_figure(ins[0], out[0]['boxes'], out[0]['labels']), it, close=True)
diff --git a/utils/data.py b/utils/data.py
index ee87c9b..63fbaf9 100644
--- a/utils/data.py
+++ b/utils/data.py
@@ -67,7 +67,7 @@ def __init__(self,
                             "Motorcyclist",
                             "Bicyclist"
                             ]
-        self.idmap = {(40,):0, (100,):1, (101,):2, (102,):3, (103,):4, (104,):5, (105,):6, (104, 41):7, (105, 41):8, (41, 104):7, (41, 105):8}
+        self.idmap = {(40,):1, (100,):2, (101,):3, (102,):4, (103,):5, (104,):6, (105,):7, (104, 41):8, (105, 41):9, (41, 104):8, (41, 105):9}
         self.cmap = np.array([
             [220, 20, 60], # person
             [  0,  0,142], # car
@@ -166,4 +166,4 @@ def to_pytorch(self, rgb, boxes, dth):
             r = Rectangle([x0, y0], x1-x0, y1-y0, alpha=.5, color=d.cmap[cl]/255.)
             axs[1].add_patch(r)
         fig.suptitle(cls)
-        plt.show()
\ No newline at end of file
+        plt.show()