hongshaorou/YOPO
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

print params first

Browse Source
This commit is contained in:
TJU_Lu 2025-07-06 21:35:32 +08:00
parent bd94cfbd51
commit b8725720da
1 changed files with 3 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
YOPO/policy/yopo_trainer.py
View File

@ -39,9 +39,6 @@ class YopoTrainer:
# params # params
self.traj_num = cfg['traj_num'] self.traj_num = cfg['traj_num']
# loss
self.yopo_loss = YOPOLoss()
# network # network
print("Loading network...") print("Loading network...")
self.policy = YopoNetwork() self.policy = YopoNetwork()
@ -53,6 +50,9 @@ class YopoTrainer:
except FileNotFoundError: except FileNotFoundError:
print("Training from scratch") print("Training from scratch")
# loss
self.yopo_loss = YOPOLoss()
# optimizer # optimizer
self.optimizer = torch.optim.AdamW(self.policy.parameters(), lr=learning_rate, fused=True) self.optimizer = torch.optim.AdamW(self.policy.parameters(), lr=learning_rate, fused=True)
print("Network Loaded! Loading Dataset...") print("Network Loaded! Loading Dataset...")