diff --git a/openicl/utils/collators.py b/openicl/utils/collators.py
index 0e6c9f6..acade3c 100644
--- a/openicl/utils/collators.py
+++ b/openicl/utils/collators.py
@@ -11,14 +11,14 @@ class ListWrapper:
     def __init__(self, data: List[Any]):
         self.data = data
 
-    def to(self, device):
+    def to(self, device, non_blocking=False):
         return self.data
 
 
 def ignore_pad_dict(features):
     res_dict = {}
     if "metadata" in features[0]:
-        res_dict['metadata'] = ListWrapper([x.pop("metadata") for x in features])
+        res_dict["metadata"] = ListWrapper([x.pop("metadata") for x in features])
     return res_dict
 
 
@@ -30,7 +30,9 @@ class DataCollatorWithPaddingAndCuda:
     max_length: Optional[int] = 3000
     pad_to_multiple_of: Optional[int] = None
 
-    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> BatchEncoding:
+    def __call__(
+        self, features: List[Dict[str, Union[List[int], torch.Tensor]]]
+    ) -> BatchEncoding:
         res_dict = ignore_pad_dict(features)
 
         has_labels = "labels" in features[0]
@@ -43,7 +45,7 @@ def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) ->
                 pad_to_multiple_of=self.pad_to_multiple_of,
                 return_attention_mask=True,
                 return_tensors="pt",
-                verbose=False
+                verbose=False,
             )
 
         # print(features)
@@ -54,11 +56,11 @@ def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) ->
             pad_to_multiple_of=self.pad_to_multiple_of,
             return_attention_mask=True,
             return_tensors="pt",
-            verbose=False
+            verbose=False,
         )
 
         if has_labels:
-            batch['labels'] = labels.input_ids
+            batch["labels"] = labels.input_ids
         batch.update(res_dict)
 
         if self.device: