GeneratorExp aren't supported in TorchScript

zhiqwang · Mar 12, 2022 · 3c48668 · 3c48668
1 parent e175764
commit 3c48668
Showing 1 changed file with 1 addition and 6 deletions.
diff --git a/yolort/models/anchor_utils.py b/yolort/models/anchor_utils.py
@@ -5,8 +5,6 @@
 import torch
 from torch import nn, Tensor
 
-from yolort.utils import check_version
-
 
 class AnchorGenerator(nn.Module):
     def __init__(self, strides: List[int], anchor_grids: List[List[float]]):
@@ -31,10 +29,7 @@ def _generate_grids(
             widths = torch.arange(width, dtype=torch.int32, device=device).to(dtype=dtype)
             heights = torch.arange(height, dtype=torch.int32, device=device).to(dtype=dtype)
 
-            if check_version(torch.__version__, "1.10.0"):
-                shift_y, shift_x = torch.meshgrid(heights, widths, indexing="ij")
-            else:
-                shift_y, shift_x = torch.meshgrid(heights, widths)
+            shift_y, shift_x = torch.meshgrid(heights, widths)
 
             grid = torch.stack((shift_x, shift_y), 2).expand((1, self.num_anchors, height, width, 2))
             grids.append(grid)