Update model.py

model.py

@@ -1587,7 +1587,7 @@ class UNet2DDragConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMi
 
         bsz, num_drags, drag_dim = drags.shape
         assert num_drags == self.num_drags
-        if (self.train and self.drag_dropout_prob > 0) or force_drop_ids is not None:
+        if (self.training and self.drag_dropout_prob > 0) or force_drop_ids is not None:
             if force_drop_ids is None:
                 drop_ids = torch.rand(bsz, device=x_cond_extra.device) < self.drag_dropout_prob
             else: