Skip to content

Commit

Permalink
Fix env.py in cases where torch is not present (#31113)
Browse files Browse the repository at this point in the history
* Fix env.py in cases where torch is not present

* Simplify the fix (and avoid some issues)
  • Loading branch information
Rocketknight1 authored May 29, 2024
1 parent c886137 commit 97a58a5
Showing 1 changed file with 8 additions and 7 deletions.
15 changes: 8 additions & 7 deletions src/transformers/commands/env.py
Original file line number Diff line number Diff line change
Expand Up @@ -133,13 +133,14 @@ def run(self):
"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))
Expand Down

0 comments on commit 97a58a5

Please sign in to comment.