torchgfn
torchgfn copied to clipboard
Use `torch.nested.nested_tensor` when it's not a prototype anymore
https://pytorch.org/docs/stable/nested.html#module-torch.nested
>>> a
tensor([0, 1, 2])
>>> b
tensor([3, 4, 5, 6, 7])
>>> nt = torch.nested.nested_tensor([a, b])
>>> nt
nested_tensor([
tensor([0, 1, 2]),
tensor([3, 4, 5, 6, 7])
])
Oh nice -- this is for storing trajectories of different lengths (for e.g.)?
Yes!