diff --git a/lib/modeling/model_builder.py b/lib/modeling/model_builder.py
index f1f5c02..e8c1279 100644
--- a/lib/modeling/model_builder.py
+++ b/lib/modeling/model_builder.py
@@ -276,7 +276,7 @@ def roi_feature_transform(self, blobs_in, rpn_ret, blob_rois='rois', method='RoI
                         # Warning!: Not check if implementation matches Detectron
                         xform_out = ROIPool((resolution, resolution), sc)(bl_in, rois)
                     elif method == 'RoIAlign':
-                        xform_out = RoIAlignFunction(
+                        xform_out = ROIAlign(
                             (resolution, resolution), sc, sampling_ratio)(bl_in, rois)
                     bl_out_list.append(xform_out)