summaryrefslogtreecommitdiff
path: root/models/model.py
diff options
context:
space:
mode:
authorJordan Gong <jordan.gong@protonmail.com>2021-02-15 14:07:20 +0800
committerJordan Gong <jordan.gong@protonmail.com>2021-02-15 14:07:20 +0800
commitbd42d8e6f3561031f957a976db10d2eb8b3d9849 (patch)
tree7433caab096d99cc7b417d332831769f5b891e7c /models/model.py
parent6a4f1cdb88d1853beb410b9da29e43069688dba7 (diff)
parent6fb1c7cb34a65769c018a08324387af419355b32 (diff)
Merge branch 'data_parallel' into data_parallel_py3.8
Diffstat (limited to 'models/model.py')
-rw-r--r--models/model.py10
1 files changed, 5 insertions, 5 deletions
diff --git a/models/model.py b/models/model.py
index 0b8de50..7bc1da9 100644
--- a/models/model.py
+++ b/models/model.py
@@ -153,10 +153,10 @@ class Model:
self.rgb_pn = nn.DataParallel(self.rgb_pn)
self.rgb_pn = self.rgb_pn.to(self.device)
self.optimizer = optim.Adam([
- {'params': self.rgb_pn.ae.parameters(), **ae_optim_hp},
- {'params': self.rgb_pn.pn.parameters(), **pn_optim_hp},
- {'params': self.rgb_pn.hpm.parameters(), **hpm_optim_hp},
- {'params': self.rgb_pn.fc_mat, **fc_optim_hp}
+ {'params': self.rgb_pn.module.ae.parameters(), **ae_optim_hp},
+ {'params': self.rgb_pn.module.pn.parameters(), **pn_optim_hp},
+ {'params': self.rgb_pn.module.hpm.parameters(), **hpm_optim_hp},
+ {'params': self.rgb_pn.module.fc_mat, **fc_optim_hp}
], **optim_hp)
sched_gamma = sched_hp.get('gamma', 0.9)
sched_step_size = sched_hp.get('step_size', 500)
@@ -195,7 +195,7 @@ class Model:
x_c2 = batch_c2['clip'].to(self.device)
y = batch_c1['label'].to(self.device)
# Duplicate labels for each part
- y = y.unsqueeze(1).repeat(1, self.rgb_pn.num_total_parts)
+ y = y.unsqueeze(1).repeat(1, self.rgb_pn.module.num_total_parts)
losses, images = self.rgb_pn(x_c1, x_c2, y)
losses = torch.stack((
# xrecon cano_cons pose_sim