Set GPU index based on local_rank instead of rank
This commit is contained in:
parent
ef98b2cef9
commit
a5f26eeef7
|
@ -227,7 +227,7 @@ def main():
|
|||
# Create and run Data Generator in a separate thread
|
||||
data_generator = DataGenerator(
|
||||
args.seed,
|
||||
hvd.rank(),
|
||||
hvd.local_rank(),
|
||||
nb_users,
|
||||
nb_items,
|
||||
neg_mat,
|
||||
|
|
Loading…
Reference in a new issue