diff --git a/models.py b/models.py
index fbd85522..5e3cd78a 100755
--- a/models.py
+++ b/models.py
@@ -254,9 +254,14 @@ class YOLOLayer(nn.Module):
                 xy = torch.sigmoid(p[..., 0:2]) + self.grid_xy  # x, y
                 width_height = torch.exp(p[..., 2:4]) * self.anchor_wh  # width, height
                 p_conf = torch.sigmoid(p[..., 4:5])  # Conf
-                ## p_cls = torch.sigmoid(p[..., 5:85])  # Class
-                p_cls = F.softmax(p[..., 5:85], 2) * p_conf  # SSD-like conf
-                # p_cls = torch.exp(p[..., 5:85]) / torch.exp(p[..., 5:85]).sum(2).unsqueeze(2) #* p_conf  # F.softmax() equivalent
+                p_cls = p[..., 5:85]
+
+                # Broadcasting only supported on first dimension in CoreML. See onnx-coreml/_operators.py
+                # p_cls = F.softmax(p_cls, 2) * p_conf  # SSD-like conf
+                p_cls = torch.exp(p_cls).permute(2, 1, 0)
+                p_cls = p_cls / p_cls.sum(0).unsqueeze(0) * p_conf.permute(2, 1, 0)  # F.softmax() equivalent
+                p_cls = p_cls.permute(2, 1, 0)
+
                 return torch.cat((xy / nG, width_height, p_conf, p_cls), 2).squeeze().t()
 
             p[..., 0] = torch.sigmoid(p[..., 0]) + self.grid_x  # x
diff --git a/utils/onnx2coreml.py b/utils/onnx2coreml.py
index f6ddecaf..2e18a763 100644
--- a/utils/onnx2coreml.py
+++ b/utils/onnx2coreml.py
@@ -68,13 +68,14 @@ def main():
         spec.neuralNetwork.preprocessing[0].featureName = '0'
 
         yolov3_model.save(name + '.mlmodel')
+        # yolov3_model.visualize_spec()
         print(spec.description)
 
         # 2.5. Try to Predict:
         from PIL import Image
         img = Image.open('../yolov3/data/samples/zidane_416.jpg')
-        out = yolov3_model.predict({'0': img})
-        print(out['141'].shape, out['143'].shape)
+        out = yolov3_model.predict({'0': img}, useCPUOnly=True)
+        print(out['148'].shape, out['150'].shape)
 
         # 3. Create NMS protobuf
         import numpy as np
@@ -106,15 +107,15 @@ def main():
             del ma_type.shape[:]
 
         nms = nms_spec.nonMaximumSuppression
-        nms.confidenceInputFeatureName = '141'  # 1x507x80
-        nms.coordinatesInputFeatureName = '143'  # 1x507x4
+        nms.confidenceInputFeatureName = '148'  # 1x507x80
+        nms.coordinatesInputFeatureName = '150'  # 1x507x4
         nms.confidenceOutputFeatureName = 'confidence'
         nms.coordinatesOutputFeatureName = 'coordinates'
         nms.iouThresholdInputFeatureName = 'iouThreshold'
         nms.confidenceThresholdInputFeatureName = 'confidenceThreshold'
 
         nms.iouThreshold = 0.6
-        nms.confidenceThreshold = 0.9
+        nms.confidenceThreshold = 0.3
         nms.pickTop.perClass = True
 
         labels = np.loadtxt('../yolov3/data/coco.names', dtype=str, delimiter='\n')