Merge pull request official-stockfish#272 from Naphthalin/fix_root_device

Sopel97 · web-flow · commit 0fa66dba8c5e · 2023-12-18T22:54:02.000+01:00
Added missing .strategy
diff --git a/train.py b/train.py
@@ -146,7 +146,7 @@ def main():
   checkpoint_callback = pl.callbacks.ModelCheckpoint(save_last=args.save_last_network, every_n_epochs=args.network_save_period, save_top_k=-1)
   trainer = pl.Trainer.from_argparse_args(args, callbacks=[checkpoint_callback], logger=tb_logger)
 
-  main_device = trainer.root_device if trainer.strategy.root_device.index is None else 'cuda:' + str(trainer.strategy.root_device.index)
+  main_device = trainer.strategy.root_device if trainer.strategy.root_device.index is None else 'cuda:' + str(trainer.strategy.root_device.index)
 
   nnue.to(device=main_device)
 
