Test with dot_product attention

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -51,7 +51,7 @@ jobs:
             train_text_encoder=false \
             cache_latents_text_encoder_outputs=true \
             per_device_batch_size=1 \
-            attention=cudnn_flash_te \
+            attention=dot_product \
             activations_dtype=bfloat16 \
             weights_dtype=bfloat16 \
             max_train_steps=200 \

src/maxdiffusion/trainers/base_stable_diffusion_trainer.py

@@ -161,8 +161,8 @@ def start_training(self):
       params["scheduler"] = noise_scheduler_state
 
       # Calculate tflops
-      # per_device_tflops = self.calculate_tflops(pipeline, params)
-      # self.per_device_tflops = per_device_tflops
+      per_device_tflops = self.calculate_tflops(pipeline, params)
+      self.per_device_tflops = per_device_tflops
 
       # Load dataset
       data_iterator = self._time_and_log_call(self.load_dataset, pipeline, params, train_states)