fix cross_attention, removed unused_parameters flag (#236)

* fix cross_attention, removed unused_paramters flag * remove kwargs handler
pinellolab · Oct 13, 2023 · 61fadc1 · 61fadc1
1 parent 02e22b4
commit 61fadc1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/src/dnadiffusion/models/unet.py b/src/dnadiffusion/models/unet.py
@@ -152,7 +152,7 @@ def forward(self, x: torch.Tensor, time: torch.Tensor, classes: torch.Tensor):
             self.norm_to_cross(x_reshaped.reshape(-1, 800)).reshape(-1, 4, 200),
             context=t_cross_reshaped,
         )  # (-1,1, 4, 200)
-        crossattention_out = x.view(-1, 1, 4, 200)
+        crossattention_out = crossattention_out.view(-1, 1, 4, 200)
         x = x + crossattention_out
         if self.output_attention:
             return x, crossattention_out

diff --git a/train.py b/train.py
@@ -7,8 +7,7 @@
 
 
 def train():
-    kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
-    accelerator = Accelerator(kwargs_handlers=[kwargs], split_batches=True, log_with=["wandb"], mixed_precision="bf16")
+    accelerator = Accelerator(split_batches=True, log_with=["wandb"], mixed_precision="bf16")
 
     data = load_data(
         data_path="src/dnadiffusion/data/K562_hESCT0_HepG2_GM12878_12k_sequences_per_group.txt",