fixes after helper refactor cleanup

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

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'
 

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):