diff --git a/examples/1_load_lerobot_dataset.py b/examples/1_load_lerobot_dataset.py index de99d69a..3846926a 100644 --- a/examples/1_load_lerobot_dataset.py +++ b/examples/1_load_lerobot_dataset.py @@ -80,7 +80,7 @@ print(f"{dataset[0]['action'].shape=}\n") # (64,c) # PyTorch datasets. dataloader = torch.utils.data.DataLoader( dataset, - num_workers=1, + num_workers=0, batch_size=32, shuffle=True, ) diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py index c5ce0d18..2d4e500b 100644 --- a/examples/3_train_policy.py +++ b/examples/3_train_policy.py @@ -49,7 +49,7 @@ optimizer = torch.optim.Adam(policy.parameters(), lr=1e-4) # Create dataloader for offline training. dataloader = torch.utils.data.DataLoader( dataset, - num_workers=4, + num_workers=0, batch_size=64, shuffle=True, pin_memory=device != torch.device("cpu"), diff --git a/tests/test_examples.py b/tests/test_examples.py index cdab31e9..9d6f97d1 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -37,7 +37,7 @@ def test_examples_3_and_2(): file_contents, [ ("training_steps = 5000", "training_steps = 1"), - ("num_workers=4", "num_workers=0"), + # ("num_workers=4", "num_workers=0"), ('device = torch.device("cuda")', 'device = torch.device("cpu")'), ("batch_size=64", "batch_size=1"), ],