Remove print statements from train.py (#101)

*Description of changes:* Removes print statements that got left inside
from a debugging session.


By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.
This commit is contained in:
Abdul Fatir 2024-06-07 23:36:45 +02:00 committed by GitHub
parent 601d52db75
commit f49c4ee4b5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -72,24 +72,17 @@ def get_training_job_info() -> Dict:
# CUDA info
job_info["cuda_available"] = torch.cuda.is_available()
print(job_info)
if torch.cuda.is_available():
job_info["device_count"] = torch.cuda.device_count()
print(job_info)
job_info["device_names"] = {
idx: torch.cuda.get_device_name(idx)
for idx in range(torch.cuda.device_count())
}
print(job_info)
print(torch.cuda.mem_get_info(device=0))
print(torch.cuda.mem_get_info(device=1))
job_info["mem_info"] = {
idx: torch.cuda.mem_get_info(device=idx)
for idx in range(torch.cuda.device_count())
}
print(job_info)
# DDP info
job_info["torchelastic_launched"] = dist.is_torchelastic_launched()