set_capture_non_tensor_stack

class tensordict.set_capture_non_tensor_stack(mode: bool)

一個上下文管理器或裝飾器，用於控制是否將相同的非張量資料堆疊到單個 NonTensorData 物件或 NonTensorStack 中。

引數:

mode (bool) – 是否捕獲非張量堆疊。如果為 False，則相同的非張量資料將被堆疊到 NonTensorStack 中。如果為 True，則單個 NonTensorData 物件將包含唯一值，但具有所需的批次大小。預設為 True。

注意

自 v0.9 起，capture_non_tensor_stack() 預設返回 False。您可以透過以下方式設定 capture_non_tensor_stack() 的值：

• 環境變數 CAPTURE_NON_TENSOR_STACK；

• 透過在指令碼開頭設定 set_capture_non_tensor_stack(val: bool).set()；

• 透過使用 set_capture_non_tensor_stack(val: bool) 作為上下文管理器或裝飾器。

建議使用 set_capture_non_tensor_stack(False) 的行為。

另請參閱

capture_non_tensor_stack

示例

>>> with set_capture_non_tensor_stack(False):
...     torch.stack([NonTensorData("a"), NonTensorData("a")])
NonTensorData("a", batch_size=[2])
>>> @set_capture_non_tensor_stack(False)
... def my_function():
...     return torch.stack([NonTensorData("a"), NonTensorData("a")])
>>> my_function()
NonTensorStack(["a", "a"], stack_dim=0)