From 3c48668ebe1bfc9415a44b1b708f730a43fe546f Mon Sep 17 00:00:00 2001 From: Zhiqiang Wang Date: Sat, 12 Mar 2022 12:53:58 +0800 Subject: [PATCH] GeneratorExp aren't supported in TorchScript --- yolort/models/anchor_utils.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/yolort/models/anchor_utils.py b/yolort/models/anchor_utils.py index 6aff8aaa..b444df59 100644 --- 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)