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
|
# Create and run Data Generator in a separate thread
|
||||||
data_generator = DataGenerator(
|
data_generator = DataGenerator(
|
||||||
args.seed,
|
args.seed,
|
||||||
hvd.rank(),
|
hvd.local_rank(),
|
||||||
nb_users,
|
nb_users,
|
||||||
nb_items,
|
nb_items,
|
||||||
neg_mat,
|
neg_mat,
|
||||||
|
|
Loading…
Reference in a new issue