diff --git a/pipeline.py b/pipeline.py
index b133012..21f9ccb 100644
--- a/pipeline.py
+++ b/pipeline.py
@@ -109,7 +109,7 @@ def encode_image_emb(self, siglip_image, device, dtype):
     def init_ccp_and_attn_processor(self, *args, **kwargs):
         subject_ip_adapter_path = kwargs['subject_ip_adapter_path']
         nb_token = kwargs['nb_token']
-        device=kwargs.get('nb_token', torch.device('cuda'))
+        device=kwargs.get('device', torch.device('cuda'))
         
         state_dict = torch.load(subject_ip_adapter_path, map_location="cpu")
         dtype = self.transformer.dtype
@@ -627,4 +627,4 @@ def enable_sequential_cpu_offload(self, gpu_id: Optional[int] = None, device: Un
                 # make sure to offload buffers if not all high level weights
                 # are of type nn.Module
                 offload_buffers = len(model._parameters) > 0
-                cpu_offload(model, device, offload_buffers=offload_buffers)
\ No newline at end of file
+                cpu_offload(model, device, offload_buffers=offload_buffers)