add log to check whether chunking is working

cosmos_transfer1/diffusion/model/model_v2w.py

@@ -249,6 +249,7 @@ class DiffusionV2WModel(DiffusionT2WModel):
         assert condition_latent is not None, "condition_latent should be provided"
 
         # try to add chunking here !!!
+        log.info("x0_fn")
         x0_fn = self.get_x0_fn_from_batch_with_condition_latent(
             data_batch,
             guidance,
@@ -312,6 +313,8 @@ class DiffusionV2WModel(DiffusionT2WModel):
             Function that takes noisy input and noise level and returns denoised prediction
         """
         if chunking is None:
+            log.info("no chunking")
+
             if is_negative_prompt:
                 condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)
             else:
@@ -347,6 +350,8 @@ class DiffusionV2WModel(DiffusionT2WModel):
 
             return x0_fn
         else:
+            log.info("chunking !!!")
+
             def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
                 if is_negative_prompt:
                     condition, uncondition = self.conditioner.get_condition_with_negative_prompt(data_batch)