diff --git a/examples/mlperf/helpers.py b/examples/mlperf/helpers.py
index dea3c20931..4929005fa2 100644
--- a/examples/mlperf/helpers.py
+++ b/examples/mlperf/helpers.py
@@ -337,7 +337,7 @@ class BoxCoder(object):
     boxes = boxes.cast(rel_codes.dtype)
     rel_codes = rel_codes
 
-    TO_REMOVE = 1  # TODO remove
+    TO_REMOVE = self.apply_to_remove  # TODO remove
     widths = boxes[:, 2] - boxes[:, 0] + TO_REMOVE
     heights = boxes[:, 3] - boxes[:, 1] + TO_REMOVE
     ctr_x = boxes[:, 0] + 0.5 * widths
diff --git a/extra/models/mask_rcnn.py b/extra/models/mask_rcnn.py
index 9937762b04..2aa813dca7 100644
--- a/extra/models/mask_rcnn.py
+++ b/extra/models/mask_rcnn.py
@@ -9,7 +9,7 @@ from tinygrad.helpers import get_child, fetch
 from tinygrad.nn.state import torch_load
 from examples.mlperf.helpers import BoxCoder
 from extra.models.resnet import ResNet
-from extra.models.helpers import nms as _box_nms
+from extra.models.retinanet import nms as _box_nms
 
 USE_NP_GATHER = os.getenv('FULL_TINYGRAD', '0') == '0'
 
diff --git a/extra/models/retinanet.py b/extra/models/retinanet.py
index 68bfd39f02..0b59da52e1 100644
--- a/extra/models/retinanet.py
+++ b/extra/models/retinanet.py
@@ -4,11 +4,10 @@ import math
 from tinygrad import Tensor, dtypes
 from tinygrad.helpers import flatten, get_child
 import tinygrad.nn as nn
-from examples.mlperf.helpers import generate_anchors
+from examples.mlperf.helpers import generate_anchors, BoxCoder
 from examples.mlperf.initializers import Conv2dNormal, Conv2dKaimingUniform
 from examples.mlperf.losses import sigmoid_focal_loss, l1_loss
 from extra.models.resnet import ResNet
-from extra.models.mask_rcnn import BoxCoder
 import numpy as np
 
 def nms(boxes, scores, thresh=0.5):