training_sem / train_hugg.py
kai-2054's picture
setup train
13dbd3b
# from huggingface_hub import HfApi
# api = HfApi()
# job = api.create_job(
# repo_id="uav-vts/training_sem",
# docker_image="pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel",
# command="python runner/train.py",
# hardware="a100",
# # token="<TOKEN>"
# )
# print("Job created:", job.job_id)
from huggingface_hub import run_job
image = "pytorch/pytorch:2.6.0-cuda12.4-cudnn9-devel"
command = [
"python",
"-c",
"import torch; print(f'This code ran with the following GPU: {torch.cuda.get_device_name()}')"
]
job = run_job(
image=image,
command=command,
flavor="a100-large",
namespace="uav-vts"
)
print("Job submitted!")
print("Track it here:", job.url)