fix: 修正 infini 使用 async h2d copy

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Jan 2, 2025 · f46d0f9 · f46d0f9
1 parent ab09e50
commit f46d0f9
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/models/llama/infini/src/lib.rs b/models/llama/infini/src/lib.rs
@@ -133,13 +133,12 @@ impl H2DLoader {
     }
 
     fn load(&mut self, host: Contiguous<HostBlob>, stream: &Stream) -> DevBlob {
-        let device = stream.get_device();
         self.event.synchronize();
         match host {
             Contiguous::Borrowed(host) => self.host.copy_from_slice(host),
             Contiguous::Owned(host) => self.host = host,
         };
-        device.memcpy_h2d(&mut self.dev, &self.host);
+        stream.memcpy_h2d(&mut self.dev, &self.host);
         stream.record(&mut self.event);
         replace(&mut self.dev, stream.malloc::<u8>(self.host.len()))
     }