Skip to content

Commit 08236be

Browse files
authored
Make readme benchmarks robust to CUDA not being present (#409)
1 parent 7e2da5d commit 08236be

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

benchmarks/decoders/generate_readme_data.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,8 @@ def main() -> None:
6161

6262
decoder_dict = {}
6363
decoder_dict["torchcodec"] = TorchCodecPublic()
64-
decoder_dict["torchcodec[cuda]"] = TorchCodecPublic(device="cuda")
64+
if torch.cuda.is_available():
65+
decoder_dict["torchcodec[cuda]"] = TorchCodecPublic(device="cuda")
6566
decoder_dict["torchvision[video_reader]"] = TorchVision("video_reader")
6667
decoder_dict["torchaudio"] = TorchAudioDecoder()
6768

@@ -80,7 +81,7 @@ def main() -> None:
8081
batch_parameters=BatchParameters(batch_size=50, num_threads=10),
8182
resize_height=256,
8283
resize_width=256,
83-
resize_device="cuda",
84+
resize_device="cuda" if torch.cuda.is_available() else "cpu",
8485
),
8586
)
8687
data_for_writing = {

0 commit comments

Comments
 (0)