This commit is contained in:
Cadene 2024-04-23 11:46:07 +00:00
parent 6729095e56
commit c29f5f3595
2 changed files with 6 additions and 6 deletions

View File

@ -51,17 +51,17 @@ def save_dataset_to_safetensors(output_dir, dataset_id="pusht"):
# TODO(rcadene): Enable testing on second and last episode # TODO(rcadene): Enable testing on second and last episode
# We currently cant because our test dataset only contains the first episode # We currently cant because our test dataset only contains the first episode
# save 2 first frames of second episode # # save 2 first frames of second episode
# i = dataset.episode_data_index["from"][1].item() # i = dataset.episode_data_index["from"][1].item()
# save_file(dataset[i], data_dir / f"frame_{i}.safetensors") # save_file(dataset[i], data_dir / f"frame_{i}.safetensors")
# save_file(dataset[i+1], data_dir / f"frame_{i+1}.safetensors") # save_file(dataset[i+1], data_dir / f"frame_{i+1}.safetensors")
# save 2 last frames of second episode # # save 2 last frames of second episode
# i = dataset.episode_data_index["to"][1].item() # i = dataset.episode_data_index["to"][1].item()
# save_file(dataset[i-2], data_dir / f"frame_{i-2}.safetensors") # save_file(dataset[i-2], data_dir / f"frame_{i-2}.safetensors")
# save_file(dataset[i-1], data_dir / f"frame_{i-1}.safetensors") # save_file(dataset[i-1], data_dir / f"frame_{i-1}.safetensors")
# save 2 last frames of last episode # # save 2 last frames of last episode
# i = dataset.episode_data_index["to"][-1].item() # i = dataset.episode_data_index["to"][-1].item()
# save_file(dataset[i-2], data_dir / f"frame_{i-2}.safetensors") # save_file(dataset[i-2], data_dir / f"frame_{i-2}.safetensors")
# save_file(dataset[i-1], data_dir / f"frame_{i-1}.safetensors") # save_file(dataset[i-1], data_dir / f"frame_{i-1}.safetensors")

View File

@ -282,17 +282,17 @@ def test_backward_compatibility():
# TODO(rcadene): Enable testing on second and last episode # TODO(rcadene): Enable testing on second and last episode
# We currently cant because our test dataset only contains the first episode # We currently cant because our test dataset only contains the first episode
# test 2 first frames of second episode # # test 2 first frames of second episode
# i = dataset.episode_data_index["from"][1].item() # i = dataset.episode_data_index["from"][1].item()
# load_and_compare(i) # load_and_compare(i)
# load_and_compare(i+1) # load_and_compare(i+1)
# test 2 last frames of second episode # #test 2 last frames of second episode
# i = dataset.episode_data_index["to"][1].item() # i = dataset.episode_data_index["to"][1].item()
# load_and_compare(i-2) # load_and_compare(i-2)
# load_and_compare(i-1) # load_and_compare(i-1)
# test 2 last frames of last episode # # test 2 last frames of last episode
# i = dataset.episode_data_index["to"][-1].item() # i = dataset.episode_data_index["to"][-1].item()
# load_and_compare(i-2) # load_and_compare(i-2)
# load_and_compare(i-1) # load_and_compare(i-1)