Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat for the GPU poors : Add GPU availability check in evaluate_pretr… #359

Merged
merged 2 commits into from
Aug 13, 2024
Merged

feat for the GPU poors : Add GPU availability check in evaluate_pretr… #359

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0c02d6c
feat for the GPU poors : Add GPU availability check in evaluate_pretr…
perezjln Aug 13, 2024
ad8e67a
Update examples/2_evaluate_pretrained_policy.py
perezjln Aug 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion examples/2_evaluate_pretrained_policy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,15 @@
output_directory = Path("outputs/eval/example_pusht_diffusion")
output_directory.mkdir(parents=True, exist_ok=True)

device = torch.device("cuda")
# Check if GPU is available
if torch.cuda.is_available():
device = torch.device("cuda")
print("GPU is available. Device set to:", device)
else:
device = torch.device("cpu")
print(f"GPU is not available. Device set to: {device}. Inference will be slower than on GPU.")
# Decrease the number of reverse-diffusion steps (trades off a bit of quality for 10x speed)
policy.diffusion.num_inference_steps = 10

# Download the diffusion policy for pusht environment
pretrained_policy_path = Path(snapshot_download("lerobot/diffusion_pusht"))
Expand Down
Loading