diff --git a/backends/vulkan/runtime/api/containers/StagingBuffer.h b/backends/vulkan/runtime/api/containers/StagingBuffer.h
index 6f67ae8a64a..1e9f569fc4a 100644
--- a/backends/vulkan/runtime/api/containers/StagingBuffer.h
+++ b/backends/vulkan/runtime/api/containers/StagingBuffer.h
@@ -23,8 +23,6 @@ class StagingBuffer final {
  private:
   Context* context_p_;
   vkapi::ScalarType dtype_;
-  size_t numel_;
-  size_t nbytes_;
   vkapi::VulkanBuffer vulkan_buffer_;
 
   void* mapped_data_;
@@ -36,10 +34,8 @@ class StagingBuffer final {
       const size_t numel)
       : context_p_(context_p),
         dtype_(dtype),
-        numel_(numel),
-        nbytes_(element_size(dtype_) * numel_),
-        vulkan_buffer_(
-            context_p_->adapter_ptr()->vma().create_staging_buffer(nbytes_)),
+        vulkan_buffer_(context_p_->adapter_ptr()->vma().create_staging_buffer(
+            element_size(dtype_) * numel)),
         mapped_data_(nullptr) {}
 
   StagingBuffer(const StagingBuffer&) = delete;
@@ -68,15 +64,15 @@ class StagingBuffer final {
   }
 
   inline size_t numel() {
-    return numel_;
+    return nbytes() / element_size(dtype_);
   }
 
   inline size_t nbytes() {
-    return nbytes_;
+    return vulkan_buffer_.mem_size();
   }
 
   inline void copy_from(const void* src, const size_t nbytes) {
-    VK_CHECK_COND(nbytes <= nbytes_);
+    VK_CHECK_COND(nbytes <= this->nbytes());
     memcpy(data(), src, nbytes);
     vmaFlushAllocation(
         vulkan_buffer_.vma_allocator(),
@@ -86,7 +82,7 @@ class StagingBuffer final {
   }
 
   inline void copy_to(void* dst, const size_t nbytes) {
-    VK_CHECK_COND(nbytes <= nbytes_);
+    VK_CHECK_COND(nbytes <= this->nbytes());
     vmaInvalidateAllocation(
         vulkan_buffer_.vma_allocator(),
         vulkan_buffer_.allocation(),
@@ -96,7 +92,7 @@ class StagingBuffer final {
   }
 
   inline void set_staging_zeros() {
-    memset(data(), 0, nbytes_);
+    memset(data(), 0, nbytes());
   }
 };