From 679708d96f1ed9da3f491eaaf7b328f7b373357b Mon Sep 17 00:00:00 2001 From: Anji Liu Date: Sun, 18 Aug 2024 21:33:25 +0800 Subject: [PATCH] relax device check in `_init_buffer` --- src/pyjuice/model/tensorcircuit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pyjuice/model/tensorcircuit.py b/src/pyjuice/model/tensorcircuit.py index 4f4d1349..7e10b5e7 100644 --- a/src/pyjuice/model/tensorcircuit.py +++ b/src/pyjuice/model/tensorcircuit.py @@ -735,7 +735,7 @@ def _init_buffer(self, name: str, shape: Tuple, set_value: Optional[float] = Non if not flag and tensor.size(i) != d: flag = True - if not flag and check_device and tensor.device != self.device: + if not flag and check_device and self.device.index is not None and tensor.device != self.device: flag = True if flag: