Fix env.py in cases where torch is not present (#31113)

* Fix env.py in cases where torch is not present

* Simplify the fix (and avoid some issues)
This commit is contained in:
Matt 2024-05-29 13:20:36 +01:00 committed by GitHub
parent c8861376ad
commit 97a58a5d2c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -133,13 +133,14 @@ class EnvironmentCommand(BaseTransformersCLICommand):
"JaxLib version": f"{jaxlib_version}",
"Using distributed or parallel set-up in script?": "<fill in>",
}
if pt_cuda_available:
info["Using GPU in script?"] = "<fill in>"
info["GPU type"] = torch.cuda.get_device_name()
elif pt_npu_available:
info["Using NPU in script?"] = "<fill in>"
info["NPU type"] = torch.npu.get_device_name()
info["CANN version"] = torch.version.cann
if is_torch_available():
if pt_cuda_available:
info["Using GPU in script?"] = "<fill in>"
info["GPU type"] = torch.cuda.get_device_name()
elif pt_npu_available:
info["Using NPU in script?"] = "<fill in>"
info["NPU type"] = torch.npu.get_device_name()
info["CANN version"] = torch.version.cann
print("\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n")
print(self.format_dict(info))