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