I am trying to build a Python library that uses Rust and WGPU to do some computation in the GPU, but the code is panicking when I create the bind group with the error `Buffer (0, 1, Vulkan) is still mapped` after the texture to buffer queue submit.

heat.wgsl

@group(0) @binding(0) var input_meta_texture: texture_3d<u32>;
@group(0) @binding(1) var input_heat_texture: texture_3d<f32>;
@group(0) @binding(2) var output_heat_texture: texture_storage_3d<r32float, write>;

@compute
@workgroup_size(4, 4, 4)
fn heat_main(
    @builtin(global_invocation_id) global_id: vec3<u32>
) {
    let dimensions = textureDimensions(input_meta_texture);
    let coords = vec3<u32>(global_id.xyz);

    if (coords.x >= dimensions.x || coords.y >= dimensions.y || coords.z >= dimensions.z) {
        return;
    }

    let current_heat = textureLoad(input_heat_texture, coords.xyz, 0);

    textureStore(output_heat_texture, coords.xyz, current_heat);
}

lib.rs

pub(crate) async fn heat_gpu(
    asset: &crate::voxels::VoxelizedAsset,
    max_amount: f32,
    iterations: usize,
    progress_callback: PyObject,
) -> Result<Vec<Vec<f32>>, crate::error::GPUError> {
    // State initializes the GPU and provides a queue to send commands to.
    let state = crate::gpu::State::new().await?;

    let (width, height, depth) = (
        asset.resolution().0 as u32,
        asset.resolution().1 as u32,
        asset.resolution().2 as u32,
    );

    let texture_size = wgpu::Extent3d {
        width,
        height,
        depth_or_array_layers: depth,
    };

    // The meta texture is a voxel grid that stores the voxel state and whether
    // or not it is a bone
    // Bit 0b01: is solid
    // Bit 0b10: is bone
    let input_meta_texture = state.device.create_texture(&wgpu::TextureDescriptor {
        label: Some("input meta texture"),
        size: texture_size,
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D3,
        format: wgpu::TextureFormat::R8Uint,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });

    // The input heat texture is a voxel grid that stores the current heat per
    // voxel.
    let input_heat_texture = state.device.create_texture(&wgpu::TextureDescriptor {
        label: Some("input heat texture"),
        size: texture_size,
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D3,
        format: wgpu::TextureFormat::R32Float,
        usage: wgpu::TextureUsages::TEXTURE_BINDING | wgpu::TextureUsages::COPY_DST,
        view_formats: &[],
    });

    // The output heat texture is a voxel grid that stores the next heat per
    // voxel.
    let output_heat_texture = state.device.create_texture(&wgpu::TextureDescriptor {
        label: Some("output heat texture"),
        size: texture_size,
        mip_level_count: 1,
        sample_count: 1,
        dimension: wgpu::TextureDimension::D3,
        format: wgpu::TextureFormat::R32Float,
        usage: wgpu::TextureUsages::TEXTURE_BINDING
            | wgpu::TextureUsages::STORAGE_BINDING
            | wgpu::TextureUsages::COPY_SRC,
        view_formats: &[],
    });

    // Set the meta texture to the voxel states (solid or not).
    let mut meta_data = asset
        .states()
        .iter()
        .map(|v| {
            if matches!(v, crate::voxels::VoxelState::Empty) {
                0
            } else {
                1
            }
        })
        .collect::<Vec<_>>();

    // Prepare the heat texture with the initial heat values.
    let mut heat_data: Vec<f32> = vec![0.0; asset.states().len()];

    let texture_bind_group_layout =
        state
            .device
            .create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: Some("heat bind group layout"),
                entries: &[
                    wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Uint,
                            view_dimension: wgpu::TextureViewDimension::D3,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 1,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::Texture {
                            sample_type: wgpu::TextureSampleType::Float { filterable: false },
                            view_dimension: wgpu::TextureViewDimension::D3,
                            multisampled: false,
                        },
                        count: None,
                    },
                    wgpu::BindGroupLayoutEntry {
                        binding: 2,
                        visibility: wgpu::ShaderStages::COMPUTE,
                        ty: wgpu::BindingType::StorageTexture {
                            access: wgpu::StorageTextureAccess::WriteOnly,
                            format: wgpu::TextureFormat::R32Float,
                            view_dimension: wgpu::TextureViewDimension::D3,
                        },
                        count: None,
                    },
                ],
            });
    let texture_bind_group = state.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("heat bind group"),
        layout: &texture_bind_group_layout,
        entries: &[
            // Input meta texture: binding 0
            wgpu::BindGroupEntry {
                binding: 0,
                resource: wgpu::BindingResource::TextureView(
                    &input_meta_texture.create_view(&wgpu::TextureViewDescriptor::default()),
                ),
            },
            // Input heat texture: binding 1
            wgpu::BindGroupEntry {
                binding: 1,
                resource: wgpu::BindingResource::TextureView(
                    &input_heat_texture.create_view(&wgpu::TextureViewDescriptor::default()),
                ),
            },
            // Output heat texture: binding 2
            wgpu::BindGroupEntry {
                binding: 2,
                resource: wgpu::BindingResource::TextureView(
                    &output_heat_texture.create_view(&wgpu::TextureViewDescriptor::default()),
                ),
            },
        ],
    });
    let shader = state
        .device
        .create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("heat shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shaders/heat.wgsl").into()),
        });

    let compute_pipeline_layout =
        state
            .device
            .create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
                label: Some("heat pipeline layout"),
                bind_group_layouts: &[&texture_bind_group_layout],
                push_constant_ranges: &[],
            });

    let pipeline = state
        .device
        .create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some("heat pipeline"),
            layout: Some(&compute_pipeline_layout),
            module: &shader,
            entry_point: "heat_main",
        });

    let (dispatch_width, dispatch_height, dispatch_depth) = crate::gpu::compute_work_group_count(
        (
            texture_size.width,
            texture_size.height,
            texture_size.depth_or_array_layers,
        ),
        crate::gpu::WORKGROUP_SIZE,
    );

    let padded_bytes_per_row =
        crate::gpu::padded_bytes_per_row(width * std::mem::size_of::<f32>() as u32);

    let output_buffer_size = padded_bytes_per_row as u64
        * height as u64
        * depth as u64
        * std::mem::size_of::<f32>() as u64;
    let output_buffer = state.device.create_buffer(&wgpu::BufferDescriptor {
        label: None,
        size: output_buffer_size,
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
        mapped_at_creation: true,
    });

    let mut output_bone_voxels = Vec::with_capacity(asset.bone_voxels().len());

    for bone_voxel in asset.bone_voxels() {
        // Prep the textures
        // The meta vec will set whether the voxel is a bone or not.
        meta_data.iter_mut().enumerate().for_each(|(i, v)| {
            if bone_voxel.voxels()[i].is_bone {
                *v |= 2
            } else {
                *v &= !2
            }
        });
        // The heat vec will set the initial heat values. Bones are max value,
        // everything else is 0.
        heat_data.iter_mut().enumerate().for_each(|(i, v)| {
            if bone_voxel.voxels()[i].is_bone {
                *v = max_amount
            } else {
                *v = 0.0
            }
        });

        state.queue.write_texture(
            input_meta_texture.as_image_copy(),
            &meta_data,
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: Some(texture_size.width * std::mem::size_of::<u8>() as u32),
                rows_per_image: Some(texture_size.height),
            },
            texture_size,
        );
        state.queue.write_texture(
            input_meta_texture.as_image_copy(),
            bytemuck::cast_slice(&heat_data),
            wgpu::ImageDataLayout {
                offset: 0,
                bytes_per_row: Some(texture_size.width * std::mem::size_of::<f32>() as u32),
                rows_per_image: Some(texture_size.height),
            },
            texture_size,
        );

        for iteration in 0..iterations {
            // Run the compute shader
            // NOTE: I believe that the encoder should be built once per
            // bone/iteration.
            let mut encoder =
                state
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some(&format!(
                            "iteration {}/{} command encoder",
                            iteration + 1,
                            iterations
                        )),
                    });

            {
                let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
                    label: Some("heat pass"),
                });
                compute_pass.set_pipeline(&pipeline);
                compute_pass.set_bind_group(0, &texture_bind_group, &[]);
                compute_pass.dispatch_workgroups(dispatch_width, dispatch_height, dispatch_depth);
            }

            // Copy the output texture to the input texture
            encoder.copy_texture_to_texture(
                wgpu::ImageCopyTexture {
                    texture: &output_heat_texture,
                    mip_level: 0,
                    origin: wgpu::Origin3d::ZERO,
                    aspect: wgpu::TextureAspect::All,
                },
                wgpu::ImageCopyTexture {
                    texture: &input_heat_texture,
                    mip_level: 0,
                    origin: wgpu::Origin3d::ZERO,
                    aspect: wgpu::TextureAspect::All,
                },
                texture_size,
            );

            state.queue.submit(Some(encoder.finish()));
            state.device.poll(wgpu::Maintain::Wait);
        }

        {
            // NOTE: I believe that the encoder should be re-built here?
            let mut encoder =
                state
                    .device
                    .create_command_encoder(&wgpu::CommandEncoderDescriptor {
                        label: Some("extract output command"),
                    });

            encoder.copy_texture_to_buffer(
                wgpu::ImageCopyTexture {
                    aspect: wgpu::TextureAspect::All,
                    texture: &output_heat_texture,
                    mip_level: 0,
                    origin: wgpu::Origin3d::ZERO,
                },
                wgpu::ImageCopyBuffer {
                    buffer: &output_buffer,
                    layout: wgpu::ImageDataLayout {
                        offset: 0,
                        bytes_per_row: Some(padded_bytes_per_row as u32),
                        rows_per_image: Some(height),
                    },
                },
                texture_size,
            );
            // Panics: thread '<unnamed>' panicked at 'Error in Queue::submit: Validation Error
            // Caused by:
            //     Buffer (0, 1, Vulkan) is still mapped
            state.queue.submit(Some(encoder.finish()));

            let buffer_slice = output_buffer.slice(..);
            buffer_slice.map_async(wgpu::MapMode::Read, |result| {
                println!("RESULT: {:?}", result);
            });

            state.device.poll(wgpu::Maintain::Wait);

            let padded_data = buffer_slice.get_mapped_range();

            let padded_heat: &[f32] = bytemuck::cast_slice(&padded_data);
            let output_bone_heat =
                padded_heat[..width as usize * height as usize * depth as usize].to_vec();
            output_bone_voxels.push(output_bone_heat);
        }

        output_buffer.unmap();
    }

    Ok(output_bone_voxels)
}