Update model.py

model.py

@@ -3,6 +3,7 @@ import tensorflow as tf
 from transformers import TFPreTrainedModel
 from tensorflow_examples.models.pix2pix import pix2pix
 from transformers import TFPreTrainedModel
+from transformers import PretrainedConfig
 
 class CycleGANConfig(PretrainedConfig):
     model_type = "cyclegan"
@@ -143,7 +144,7 @@ class TFCycleGANModel(TFPreTrainedModel):
           gen_g_loss = self.generator_loss(disc_fake_y)
           gen_f_loss = self.generator_loss(disc_fake_x)
       
-          total_cycle_loss = self.calc_cycle_loss(real_x, cycled_x) + calc_cycle_loss(real_y, cycled_y)
+          total_cycle_loss = self.calc_cycle_loss(real_x, cycled_x) + self.calc_cycle_loss(real_y, cycled_y)
       
           # # Total generator loss = adversarial loss + cycle loss
           total_gen_g_loss = gen_g_loss + total_cycle_loss + self.identity_loss(real_y, same_y)
@@ -157,27 +158,27 @@ class TFCycleGANModel(TFPreTrainedModel):
       
         # Calculate the gradients for generator and discriminator
         generator_g_gradients = tape.gradient(total_gen_g_loss,
-                                              generator_g.trainable_variables)
+                                              self.generator_g.trainable_variables)
         generator_f_gradients = tape.gradient(total_gen_f_loss,
-                                              generator_f.trainable_variables)
+                                              self.generator_f.trainable_variables)
       
         discriminator_x_gradients = tape.gradient(disc_x_loss,
-                                                  discriminator_x.trainable_variables)
+                                                  self.discriminator_x.trainable_variables)
         discriminator_y_gradients = tape.gradient(disc_y_loss,
-                                                  discriminator_y.trainable_variables)
+                                                  self.discriminator_y.trainable_variables)
       
         # Apply the gradients to the optimizer
-        generator_g_optimizer.apply_gradients(zip(generator_g_gradients,
-                                                  generator_g.trainable_variables))
+        self.generator_g_optimizer.apply_gradients(zip(generator_g_gradients,
+                                                  self.generator_g.trainable_variables))
       
-        generator_f_optimizer.apply_gradients(zip(generator_f_gradients,
-                                                  generator_f.trainable_variables))
+        self.generator_f_optimizer.apply_gradients(zip(generator_f_gradients,
+                                                  self.generator_f.trainable_variables))
       
-        discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients,
-                                                      discriminator_x.trainable_variables))
+        self.discriminator_x_optimizer.apply_gradients(zip(discriminator_x_gradients,
+                                                      self.discriminator_x.trainable_variables))
       
-        discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients,
-                                                      discriminator_y.trainable_variables))
+        self.discriminator_y_optimizer.apply_gradients(zip(discriminator_y_gradients,
+                                                      self.discriminator_y.trainable_variables))