Commit 1d9b0a99 authored by chicm-ms's avatar chicm-ms Committed by SparkSnail
Browse files

Fix local gpu issue (#1000)

parent b7ca02c7
......@@ -369,7 +369,7 @@ class LocalTrainingService implements TrainingService {
private getEnvironmentVariables(
trialJobDetail: TrialJobDetail,
resource?: { gpuIndices: number[] }): { key: string; value: string }[] {
resource: { gpuIndices: number[] }): { key: string; value: string }[] {
const envVariables: { key: string; value: string }[] = [
{ key: 'NNI_PLATFORM', value: 'local' },
{ key: 'NNI_SYS_DIR', value: trialJobDetail.workingDirectory },
......@@ -379,12 +379,10 @@ class LocalTrainingService implements TrainingService {
{ key: 'MULTI_PHASE', value: this.isMultiPhase.toString() }
];
if (resource !== undefined && resource.gpuIndices.length > 0) {
envVariables.push({
key: 'CUDA_VISIBLE_DEVICES',
value: this.gpuScheduler === undefined ? '' : resource.gpuIndices.join(',')
});
}
return envVariables;
}
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment