Unverified Commit b0e8ece9 authored by LRJKD's avatar LRJKD Committed by GitHub
Browse files

[Fix] Fix support for devices other than Cuda in FreeAnchor3DHead (#2769)

parent 0f9dfa97
...@@ -2,6 +2,7 @@ ...@@ -2,6 +2,7 @@
from typing import Dict, List from typing import Dict, List
import torch import torch
from mmengine.device import get_device
from torch import Tensor from torch import Tensor
from torch.nn import functional as F from torch.nn import functional as F
...@@ -79,7 +80,9 @@ class FreeAnchor3DHead(Anchor3DHead): ...@@ -79,7 +80,9 @@ class FreeAnchor3DHead(Anchor3DHead):
featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores] featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
assert len(featmap_sizes) == self.prior_generator.num_levels assert len(featmap_sizes) == self.prior_generator.num_levels
anchor_list = self.get_anchors(featmap_sizes, batch_input_metas) device = get_device()
anchor_list = self.get_anchors(featmap_sizes, batch_input_metas,
device)
mlvl_anchors = [torch.cat(anchor) for anchor in anchor_list] mlvl_anchors = [torch.cat(anchor) for anchor in anchor_list]
# concatenate each level # concatenate each level
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment