[Improve] Add exception for PointRend for support CPU-only usage (#1271)
* [Improve] Add exception for PointRend for support CPU-only usage * fixed lintingpull/1801/head
parent
548eec9580
commit
b405ce0506
|
@ -4,7 +4,11 @@
|
|||
import torch
|
||||
import torch.nn as nn
|
||||
from mmcv.cnn import ConvModule
|
||||
from mmcv.ops import point_sample
|
||||
|
||||
try:
|
||||
from mmcv.ops import point_sample
|
||||
except ModuleNotFoundError:
|
||||
point_sample = None
|
||||
|
||||
from mmseg.models.builder import HEADS
|
||||
from mmseg.ops import resize
|
||||
|
@ -75,6 +79,9 @@ class PointHead(BaseCascadeDecodeHead):
|
|||
init_cfg=dict(
|
||||
type='Normal', std=0.01, override=dict(name='fc_seg')),
|
||||
**kwargs)
|
||||
if point_sample is None:
|
||||
raise RuntimeError('Please install mmcv-full for '
|
||||
'point_sample ops')
|
||||
|
||||
self.num_fcs = num_fcs
|
||||
self.coarse_pred_each_layer = coarse_pred_each_layer
|
||||
|
|
Loading…
Reference in New Issue