Remove In-place operations

modeling_molmo.py

@@ -1163,7 +1163,7 @@ class MultiHeadAttentionPool(nn.Module):
         if self.dropout:
             attn_output = self.residual_dropout(attn_output)
         if self.mean_residual:
-            attn_output += inputs_kv.mean(dim=1, keepdim=True)
+            attn_output = attn_output + inputs_kv.mean(dim=1, keepdim=True)
 
         return attn_output
 
@@ -1879,7 +1879,7 @@ class Molmo(nn.Module):
             # For hf demo/endpoint
             image_features = image_features.to(x.device)
 
-            x[batch_idx[valid], image_input_idx[valid]] += image_features[valid]
+            x[batch_idx[valid], image_input_idx[valid]] = x[batch_idx[valid], image_input_idx[valid]] + image_features[valid]
 
         if not self.config.rope:
             # Get positional embeddings.
@@ -2145,7 +2145,7 @@ class MolmoForCausalLM(PreTrainedModel):
                     z_loss = z_loss.view(input_ids.shape[0], -1)
                     z_loss = z_loss * loss_masks
                     z_loss = z_loss.sum() / batch_size_in_tokens
-                    loss += z_loss
+                    loss = loss + z_loss
             else:
                 # Shift so that tokens < n predict n
                 shift_logits = logits[..., :-1, :].contiguous()