RoyYang0714 commited on
Commit
048731a
·
1 Parent(s): 1f827e2

fix: Handle cuda issue.

Browse files
Files changed (2) hide show
  1. app.py +1 -2
  2. vis4d/op/layer/ms_deform_attn.py +16 -8
app.py CHANGED
@@ -86,8 +86,7 @@ def run_3d_mood(fx, fy, cx, cy, image):
86
 
87
  gc.collect()
88
 
89
- # device = "cuda" if torch.cuda.is_available() else "cpu"
90
- device = "cpu"
91
 
92
  # Data
93
  images = image.astype(np.float32)[None, ...]
 
86
 
87
  gc.collect()
88
 
89
+ device = "cuda" if torch.cuda.is_available() else "cpu"
 
90
 
91
  # Data
92
  images = image.astype(np.float32)[None, ...]
vis4d/op/layer/ms_deform_attn.py CHANGED
@@ -542,14 +542,22 @@ class MultiScaleDeformableAttention(nn.Module):
542
  )
543
 
544
  if torch.cuda.is_available() and value.is_cuda:
545
- output = MSDeformAttentionFunction.apply(
546
- value,
547
- input_spatial_shapes,
548
- input_level_start_index,
549
- sampling_locations,
550
- attention_weights,
551
- self.im2col_step,
552
- )
 
 
 
 
 
 
 
 
553
  else:
554
  output = ms_deformable_attention_cpu(
555
  value,
 
542
  )
543
 
544
  if torch.cuda.is_available() and value.is_cuda:
545
+ if VIS4D_CUDA_OPS_AVAILABLE:
546
+ output = MSDeformAttentionFunction.apply(
547
+ value,
548
+ input_spatial_shapes,
549
+ input_level_start_index,
550
+ sampling_locations,
551
+ attention_weights,
552
+ self.im2col_step,
553
+ )
554
+ else:
555
+ output = ms_deformable_attention_cpu(
556
+ value.cpu(),
557
+ input_spatial_shapes.cpu(),
558
+ sampling_locations.cpu(),
559
+ attention_weights.cpu(),
560
+ ).cuda()
561
  else:
562
  output = ms_deformable_attention_cpu(
563
  value,