Skip to content
Permalink
Browse files

fixed gpu selection bug; thanks Arnab!

  • Loading branch information...
srijankr
srijankr committed Jul 23, 2019
1 parent 51f06b9 commit 0f575ab98130c4c48e4e70fc3bcc0a330107366c
Showing with 1 addition and 1 deletion.
  1. +1 −1 library_models.py
@@ -216,7 +216,7 @@ def set_embeddings_training_end(user_embeddings, item_embeddings, user_embedding
# SELECT THE GPU WITH MOST FREE MEMORY TO SCHEDULE JOB
def select_free_gpu():
mem = []
gpus = list(set([0,1,2,3]))
gpus = list(set(range(torch.cuda.device_count()))) # list(set(X)) is done to shuffle the array
for i in gpus:
gpu_stats = gpustat.GPUStatCollection.new_query()
mem.append(gpu_stats.jsonify()["gpus"][i]["memory.used"])

0 comments on commit 0f575ab

Please sign in to comment.
You can’t perform that action at this time.