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

incorporated change from PR #10

Browse Source
This commit is contained in:
Cheng Chi 2023-09-07 19:02:13 -04:00
parent 0d00e02b45
commit dd2cbac9fa
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
diffusion_policy/model/diffusion/conditional_unet1d.py
View File

@ -226,6 +226,10 @@ class ConditionalUnet1D(nn.Module):
for idx, (resnet, resnet2, upsample) in enumerate(self.up_modules):
x = torch.cat((x, h.pop()), dim=1)
x = resnet(x, global_feature)
# The correct condition should be:
# if idx == (len(self.up_modules)-1) and len(h_local) > 0:
# However this change will break compatibility with published checkpoints.
# Therefore it is left as a comment.
if idx == len(self.up_modules) and len(h_local) > 0:
x = x + h_local[1]
x = resnet2(x, global_feature)