Fix custom imports for both distributed and single device #1731
Dr.CI classification results
{"FAILED":[],"FLAKY":[],"BROKEN_TRUNK":[{"workflowId":11133416055,"workflowUniqueId":103397641,"id":30942077270,"runnerName":"i-042bd7c5e3760c6ff","authorEmail":"rafiayub@fb.com","name":"GPU tests / gpu_test (3.9, stable)","jobName":"gpu_test (3.9, stable)","conclusion":"cancelled","completed_at":"2024-10-01T22:12:46.000000000Z","html_url":"https://github.com/pytorch/torchtune/actions/runs/11133416055/job/30942077270","head_branch":"fix_custom_imports","pr_number":1731,"head_sha":"d7be38fa2dd00cc6b92b141e95e5c41ede5a866a","head_sha_timestamp":"2024-10-01T21:02:58Z","failure_captures":["##[error]The operation was canceled."],"failure_lines":["##[error]The operation was canceled."],"failure_context":[],"time":"2024-10-01T22:10:59.000000000Z"}],"UNSTABLE":[]}