Merge 0c91dea4a2 into a3c7b01cc0
commit
f1adf75857
|
|
@ -476,7 +476,6 @@ if __name__ == "__main__":
|
|||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger('torch.compile enabled')
|
||||
rollout_engine.update_policy(model)
|
||||
if dist.is_initialized():
|
||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||
rollout_engine.update_policy(model)
|
||||
|
|
|
|||
|
|
@ -311,7 +311,6 @@ if __name__ == "__main__":
|
|||
if args.use_compile == 1:
|
||||
model = torch.compile(model)
|
||||
Logger('torch.compile enabled')
|
||||
rollout_engine.update_policy(model)
|
||||
if dist.is_initialized():
|
||||
model = DistributedDataParallel(model, device_ids=[local_rank])
|
||||
rollout_engine.update_policy(model)
|
||||
|
|
|
|||
|
|
@ -427,7 +427,6 @@ if __name__ == "__main__":
|
|||
if args.use_compile == 1:
|
||||
actor_model = torch.compile(actor_model)
|
||||
Logger('torch.compile enabled')
|
||||
rollout_engine.update_policy(actor_model)
|
||||
if dist.is_initialized():
|
||||
actor_model = DistributedDataParallel(actor_model, device_ids=[local_rank])
|
||||
critic_model = DistributedDataParallel(critic_model, device_ids=[local_rank])
|
||||
|
|
|
|||
Loading…
Reference in New Issue