Update modeling-dev2.py

modeling-dev2.py

@@ -126,25 +126,31 @@ class MiniSunModel(tf.keras.Model):
         inputs, labels = data
         input_ids = inputs['input_ids']
         attention_mask = inputs['attention_mask']
-
+    
         with tf.GradientTape() as tape:
             logits, _ = self(inputs, training=True)
             loss = self.compute_loss(labels, logits)
-
+    
         gradients = tape.gradient(loss, self.trainable_variables)
-
+    
         # Gradient clipping for stability
         gradients = [tf.clip_by_value(g, -1.0, 1.0) for g in gradients]
         self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
-
+    
+        # Compute predictions and metrics
         logits_for_metrics = tf.argmax(logits, axis=-1)
-        logits_for_metrics = tf.reshape(logits_for_metrics, [-1])
-        labels_for_metrics = tf.reshape(labels, [-1])
-
+        labels_for_metrics = tf.reshape(labels, [-1])  # Flatten labels
+        logits_for_metrics = tf.reshape(logits_for_metrics, [-1])  # Flatten predictions
+    
         for metric in self.metrics:
             metric.update_state(labels_for_metrics, logits_for_metrics)
+    
+        # Return loss and metrics
+        results = {m.name: m.result() for m in self.metrics}
+        results['loss'] = loss
+        
+        return results
 
-        return {m.name: m.result() for m in self.metrics}
 
 def create_model(config):
     model = MiniSunModel(config)