From 764a7ad2c3b98ba1f2e968ee9a6eaa0fc8be8f8b Mon Sep 17 00:00:00 2001 From: Cadene Date: Mon, 6 May 2024 00:49:33 +0000 Subject: [PATCH] try --- examples/1_load_lerobot_dataset.py | 2 +- examples/3_train_policy.py | 2 +- tests/test_examples.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) 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"), ],