fix bug of batch_size omission

Author: baymax591

verl/protocol.py

@@ -198,12 +198,16 @@ def union_numpy_dict(tensor_dict1: dict[str, np.ndarray], tensor_dict2: dict[str
     return tensor_dict1
 
 
-def numpy_dict_to_tensor_dict(numpy_dict: dict[str, np.ndarray]) -> TensorDict:
+def numpy_dict_to_tensor_dict(numpy_dict: dict[str, np.ndarray], batch_size=None) -> TensorDict:
     """Convert a dictionary of numpy arrays to a tensordict"""
-    tensor_dict = tensordict.TensorDict()
+    tensor_dict = {}
     for key, val in numpy_dict.items():
-        tensor_dict[key] = torch.from_numpy(val)
-    return tensor_dict
+        if isinstance(val, np.ndarray):
+            tensor_dict[key] = torch.from_numpy(val)
+        else:
+            tensor_dict[key] = val
+    
+    return TensorDict(tensor_dict, batch_size=batch_size)
 
 
 def list_of_dict_to_dict_of_list(list_of_dict: list[dict]):
@@ -339,18 +343,13 @@ def __getitem__(self, item):
             raise TypeError(f"Indexing with {type(item)} is not supported")
 
     def __getstate__(self):
-        return pickle.dumps(self.batch.numpy()), self.non_tensor_batch, self.meta_info
+        return pickle.dumps(self.batch.numpy()), self.batch.batch_size, self.non_tensor_batch, self.meta_info
 
     def __setstate__(self, data):
-        batch_deserialized_bytes, non_tensor_batch, meta_info = data
+        batch_deserialized_bytes, batch_size, non_tensor_batch, meta_info = data
         batch_deserialized = pickle.loads(batch_deserialized_bytes)
-
-        tensor_dict = torch.utils._pytree.tree_map(
-            lambda x: torch.from_numpy(x) if isinstance(x, np.ndarray) else x,
-            batch_deserialized
-        )
-
-        self.batch = TensorDict.from_dict(tensor_dict)
+        
+        self.batch = numpy_dict_to_tensor_dict(batch_deserialized, batch_size=batch_size)
         self.non_tensor_batch = non_tensor_batch
         self.meta_info = meta_info