search

pytorch / tensordict

TensorDict is a pytorch dedicated tensor container.
MIT License
832 stars 74 forks source link

[Feature] Compile - benchmarks #883

Closed vmoens closed 3 months ago

vmoens commented 3 months ago

Stack from ghstack (oldest at bottom):

vmoens commented 3 months ago

@jisaacso this is the last PR of the compile stack

It includes benchmarks for tensordict vs pytree in eager and compile. Ops like copy (test: test_compile_copy_flat) run faster in eager and compile with TD compared to pytree.