fix comment

Author: zhabuye

recipe/simple_use_case/async_demo.py

@@ -66,7 +66,7 @@ def actor_rollout_wg_generate_sequences(self, data_meta, data_system_client):
         output = TensorDict(
             {
                 "generate_sequences_ids": output,
-                "non_tensor_data": torch.stack([NonTensorData("test_str") for _ in range(output.size(0))]),
+                "non_tensor_data": NonTensorData(["test_str" for _ in range(output.size(0))]),
                 "nested_tensor": torch.nested.as_nested_tensor([torch.randn(1, 2) for _ in range(output.size(0))]),
             },
             batch_size=output.size(0),
@@ -118,7 +118,7 @@ async def generate(self, data_meta):
         output = TensorDict(
             {
                 "generate_sequences_ids": data,
-                "non_tensor_data": torch.stack([NonTensorData("test_str") for _ in range(data.size(0))]),
+                "non_tensor_data": NonTensorData(["test_str" for _ in range(data.size(0))]),
                 "nested_tensor": torch.nested.as_nested_tensor([torch.randn(1, 2) for _ in range(data.size(0))]),
             },
             batch_size=data.size(0),

transfer_queue/storage/simple_backend.py

@@ -86,6 +86,8 @@ def get_data(self, fields: list[str], local_indexes: list[int]) -> TensorDict[st
             if len(local_indexes) == 1:
                 # The unsqueeze op make the shape from n to (1, n)
                 gathered_item = self.field_data[field][local_indexes[0]]
+                if gathered_item is None:
+                    raise ValueError(f"Missing data for field '{field}' at index {local_indexes[0]}")
                 if not isinstance(gathered_item, torch.Tensor):
                     result[field] = NonTensorStack(gathered_item)
                 else:
@@ -94,13 +96,18 @@ def get_data(self, fields: list[str], local_indexes: list[int]) -> TensorDict[st
                 gathered_items = list(itemgetter(*local_indexes)(self.field_data[field]))
 
                 if gathered_items:
+                    if any(x is None for x in gathered_items):
+                        missing = [i for i, x in zip(local_indexes, gathered_items) if x is None]
+                        raise ValueError(f"Missing data for field '{field}' at indexes {missing}")
                     all_tensors = all(isinstance(x, torch.Tensor) for x in gathered_items)
                     if all_tensors:
                         result[field] = torch.nested.as_nested_tensor(gathered_items)
                     else:
                         result[field] = NonTensorStack(*gathered_items)
 
-        return TensorDict(result)
+        # Explicit batch size for stability
+        bs = 0 if not fields or not local_indexes else len(local_indexes)
+        return TensorDict(result, batch_size=bs)
 
     def put_data(self, field_data: TensorDict[str, Any], local_indexes: list[int]) -> None:
         """
@@ -110,7 +117,15 @@ def put_data(self, field_data: TensorDict[str, Any], local_indexes: list[int]) -
             field_data: Dict with field names as keys, corresponding data in the field as values.
             local_indexes: Local indexes used for putting data.
         """
-        extracted_data = dict(field_data)
+        # Accept TensorDict or plain dict[str, list-like]
+        if isinstance(field_data, TensorDict):
+            extracted_data = field_data.to_dict()
+        elif isinstance(field_data, dict):
+            extracted_data = field_data
+        else:
+            raise TypeError(
+                f"field_data must be a TensorDict or dict[str, list-like], got {type(field_data)}"
+            )
 
         for f, values in extracted_data.items():
             if f not in self.field_data: