Update

deepjavalibrary · Aug 16, 2024 · 05886d4 · 05886d4
1 parent 8cb82fd
commit 05886d4
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsModel.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsModel.java
@@ -28,7 +28,6 @@
 /** {@code RsModel} is the Rust implementation of {@link Model}. */
 public class RsModel extends BaseModel {
 
-    private Device device;
     private final AtomicReference<Long> handle;
 
     /**
@@ -39,7 +38,6 @@ public class RsModel extends BaseModel {
      */
     RsModel(String name, Device device) {
         super(name);
-        this.device = device;
         manager = RsNDManager.getSystemManager().newSubManager(device);
         manager.setName("RsModel");
         dataType = DataType.FLOAT16;
@@ -60,8 +58,8 @@ public void load(Path modelPath, String prefix, Map<String, ?> options)
                     RustLibrary.loadModel(
                             modelDir.toAbsolutePath().toString(),
                             dataType.ordinal(),
-                            device.getDeviceType(),
-                            device.getDeviceId()));
+                            manager.getDevice().getDeviceType(),
+                            manager.getDevice().getDeviceId()));
             block = new RsSymbolBlock((RsNDManager) manager, handle.get());
         } else {
             loadBlock(prefix, options);