Atttempt to reduce unnecessary cuda sync (#5515)

xuzhao9 · datumbox · web-flow · commit a784db48e84b · 2022-03-04T11:01:22.000Z
Co-authored-by: Vasilis Vryniotis &lt;datumbox@users.noreply.github.com&gt;
diff --git a/torchvision/models/detection/anchor_utils.py b/torchvision/models/detection/anchor_utils.py
@@ -121,8 +121,8 @@ def forward(self, image_list: ImageList, feature_maps: List[Tensor]) -> List[Ten
         dtype, device = feature_maps[0].dtype, feature_maps[0].device
         strides = [
             [
-                torch.tensor(image_size[0] // g[0], dtype=torch.int64, device=device),
-                torch.tensor(image_size[1] // g[1], dtype=torch.int64, device=device),
+                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[0] // g[0]),
+                torch.empty((), dtype=torch.int64, device=device).fill_(image_size[1] // g[1]),
             ]
             for g in grid_sizes
         ]

Original file line number	Diff line number	Diff line change
`@@ -121,8 +121,8 @@ def forward(self, image_list: ImageList, feature_maps: List[Tensor]) -> List[Ten`
`121`	`121`	`dtype, device = feature_maps[0].dtype, feature_maps[0].device`
`122`	`122`	`strides = [`
`123`	`123`	`[`
`124`		`- torch.tensor(image_size[0] // g[0], dtype=torch.int64, device=device),`
`125`		`- torch.tensor(image_size[1] // g[1], dtype=torch.int64, device=device),`
	`124`	`+ torch.empty((), dtype=torch.int64, device=device).fill_(image_size[0] // g[0]),`
	`125`	`+ torch.empty((), dtype=torch.int64, device=device).fill_(image_size[1] // g[1]),`
`126`	`126`	`]`
`127`	`127`	`for g in grid_sizes`
`128`	`128`	`]`