Spaces:
Running
on
Zero
Running
on
Zero
Commit
·
048731a
1
Parent(s):
1f827e2
fix: Handle cuda issue.
Browse files- app.py +1 -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 |
-
|
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 |
-
|
546 |
-
|
547 |
-
|
548 |
-
|
549 |
-
|
550 |
-
|
551 |
-
|
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,
|