updated README.

README.md

@@ -18,7 +18,7 @@ model-index:
           link: https://www.kaggle.com/datasets/mateuszbuda/lgg-mri-segmentation
         metrics:
           - type: accuracy
-            value: 0.958
+            value: 0.953
             name: Accuracy
 ---
 
@@ -60,7 +60,7 @@ input_data = tf.ones((1, 256, 256, 3), dtype=tf.float32)
 output = model(input_data)["output_0"]
 
 # Prints the predicted class (e.g., 0 for normal, 1 for abnormal)
-predicted_class = tf.argmax(output, axis=-1).numpy()
+predicted_class = tf.argmax(output, axis=-1).numpy()[0]
 print("Predicted class: ", predicted_class)
 ```
 
@@ -69,7 +69,7 @@ print("Predicted class: ", predicted_class)
 ### Compute
 
 - The model was trained on a GeForce 4070Ti GPU with 16GB VRAM.
-- Training completed in approximately 1.8 minutes over 20 epochs.
+- Training completed in approximately 1.6 minutes over 17 epochs.
 
 ### Dataset
 
@@ -77,7 +77,7 @@ print("Predicted class: ", predicted_class)
 
 ### Performance on test set
 
-- **Accuracy:** 0.958
+- **Accuracy:** 0.953
 
 ## Citation