From 6c4a883a9e67330fdecb0982b0c0611d425c8681 Mon Sep 17 00:00:00 2001 From: Jianzhu Guo Date: Sun, 2 Feb 2025 17:06:01 +0800 Subject: [PATCH] fix: issue 483 (#484) --- src/utils/animal_landmark_runner.py | 2 +- .../UniPose/ops/src/cuda/ms_deform_attn_cuda.cu | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/utils/animal_landmark_runner.py b/src/utils/animal_landmark_runner.py index dd91aa5..c203004 100644 --- a/src/utils/animal_landmark_runner.py +++ b/src/utils/animal_landmark_runner.py @@ -43,7 +43,7 @@ class XPoseRunner(object): args = Config.fromfile(model_config_path) args.device = device model = build_model(args) - checkpoint = torch.load(model_checkpoint_path, map_location=lambda storage, loc: storage) + checkpoint = torch.load(model_checkpoint_path, map_location=lambda storage, loc: storage, weights_only=False) load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False) model.eval() return model diff --git a/src/utils/dependencies/XPose/models/UniPose/ops/src/cuda/ms_deform_attn_cuda.cu b/src/utils/dependencies/XPose/models/UniPose/ops/src/cuda/ms_deform_attn_cuda.cu index d6d5836..913a5fc 100644 --- a/src/utils/dependencies/XPose/models/UniPose/ops/src/cuda/ms_deform_attn_cuda.cu +++ b/src/utils/dependencies/XPose/models/UniPose/ops/src/cuda/ms_deform_attn_cuda.cu @@ -18,7 +18,7 @@ at::Tensor ms_deform_attn_cuda_forward( - const at::Tensor &value, + const at::Tensor &value, const at::Tensor &spatial_shapes, const at::Tensor &level_start_index, const at::Tensor &sampling_loc, @@ -50,7 +50,7 @@ at::Tensor ms_deform_attn_cuda_forward( const int im2col_step_ = std::min(batch, im2col_step); AT_ASSERTM(batch % im2col_step_ == 0, "batch(%d) must divide im2col_step(%d)", batch, im2col_step_); - + auto output = at::zeros({batch, num_query, num_heads, channels}, value.options()); const int batch_n = im2col_step_; @@ -61,7 +61,7 @@ at::Tensor ms_deform_attn_cuda_forward( for (int n = 0; n < batch/im2col_step_; ++n) { auto columns = output_n.select(0, n); - AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_forward_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_forward_cuda", ([&] { ms_deformable_im2col_cuda(at::cuda::getCurrentCUDAStream(), value.data() + n * im2col_step_ * per_value_size, spatial_shapes.data(), @@ -81,7 +81,7 @@ at::Tensor ms_deform_attn_cuda_forward( std::vector ms_deform_attn_cuda_backward( - const at::Tensor &value, + const at::Tensor &value, const at::Tensor &spatial_shapes, const at::Tensor &level_start_index, const at::Tensor &sampling_loc, @@ -127,11 +127,11 @@ std::vector ms_deform_attn_cuda_backward( auto per_sample_loc_size = num_query * num_heads * num_levels * num_point * 2; auto per_attn_weight_size = num_query * num_heads * num_levels * num_point; auto grad_output_n = grad_output.view({batch/im2col_step_, batch_n, num_query, num_heads, channels}); - + for (int n = 0; n < batch/im2col_step_; ++n) { auto grad_output_g = grad_output_n.select(0, n); - AT_DISPATCH_FLOATING_TYPES(value.type(), "ms_deform_attn_backward_cuda", ([&] { + AT_DISPATCH_FLOATING_TYPES(value.scalar_type(), "ms_deform_attn_backward_cuda", ([&] { ms_deformable_col2im_cuda(at::cuda::getCurrentCUDAStream(), grad_output_g.data(), value.data() + n * im2col_step_ * per_value_size, @@ -150,4 +150,4 @@ std::vector ms_deform_attn_cuda_backward( return { grad_value, grad_sampling_loc, grad_attn_weight }; -} \ No newline at end of file +}