From 0fa96b96d33a944186b129deffcf32fd3181281a Mon Sep 17 00:00:00 2001
From: Xiaochun Tong <xtong@uwaterloo.ca>
Date: Sun, 24 Sep 2023 02:10:46 -0400
Subject: [PATCH] minor fix

---
 luisa_compute/src/lang.rs    | 4 ++--
 luisa_compute/src/runtime.rs | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/luisa_compute/src/lang.rs b/luisa_compute/src/lang.rs
index a6fada1..8e22085 100644
--- a/luisa_compute/src/lang.rs
+++ b/luisa_compute/src/lang.rs
@@ -282,7 +282,7 @@ impl Recorder {
         self.callable_ret_type = None;
     }
 
-    pub(crate) fn check_on_same_device(&self, other: &Device) -> Option<(String, String)> {
+    pub(crate) fn check_on_same_device(&mut self, other: &Device) -> Option<(String, String)> {
         if let Some(device) = &self.device {
             let device = device.upgrade().unwrap();
             if !Arc::ptr_eq(&device.inner, &other.inner) {
@@ -292,7 +292,7 @@ impl Recorder {
                 ));
             }
         } else {
-            // @FIXME: What should we do?
+            self.device = Some(WeakDevice::new(other));
         }
         None
     }
diff --git a/luisa_compute/src/runtime.rs b/luisa_compute/src/runtime.rs
index a1409b1..f83aecf 100644
--- a/luisa_compute/src/runtime.rs
+++ b/luisa_compute/src/runtime.rs
@@ -1254,7 +1254,7 @@ pub struct RawCallable {
 impl RawCallable {
     pub(crate) fn check_on_same_device(&self) {
         RECORDER.with(|r| {
-            let r =r.borrow();
+            let mut r =r.borrow_mut();
             if let Some(device) = &self.device {
                 if let Some((a,b)) = r.check_on_same_device(device) {
                     panic!("Callable created on a different device than the one it is called on: {:?} vs {:?}", a,b);