Comment out tflop calc

Author: Rohan-Bierneni

.github/workflows/UnitTests.yml

@@ -37,8 +37,8 @@ jobs:
 
       - name: Print dependencies
         run: |
-          pip uninstall -y transformer-engine transformer-engine-jax
-          pip install -U transformer-engine[pytorch,jax]
+          # pip uninstall -y transformer-engine transformer-engine-jax
+          # pip install -U transformer-engine[pytorch,jax]
           pip freeze
 
       - name: Run MaxDiffusion Training

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)