Commit 9f2ae1cb authored by KohakuBlueleaf's avatar KohakuBlueleaf

Add missing .mean

parent 3e0146f9
...@@ -173,7 +173,7 @@ class NetworkModule: ...@@ -173,7 +173,7 @@ class NetworkModule:
orig_weight = orig_weight.to(updown) orig_weight = orig_weight.to(updown)
merged_scale1 = updown + orig_weight merged_scale1 = updown + orig_weight
dora_merged = ( dora_merged = (
merged_scale1 / merged_scale1(dim=self.dora_mean_dim, keepdim=True) * self.dora_scale merged_scale1 / merged_scale1.mean(dim=self.dora_mean_dim, keepdim=True) * self.dora_scale
) )
final_updown = dora_merged - orig_weight final_updown = dora_merged - orig_weight
return final_updown return final_updown
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment