diff --git a/skin_cancer_classification.py b/skin_cancer_classification.py
index 7d38c2f..11c2ecd 100644
--- a/skin_cancer_classification.py
+++ b/skin_cancer_classification.py
@@ -29,9 +29,9 @@ Original file is located at
 import sys
 IN_COLAB = 'google.colab' in sys.modules
 
-#if IN_COLAB:
-#  !pip install pandas numpy matplotlib seaborn pillow scikit-learn tensorflow
-#  !pip install --upgrade kagglehub[pandas-datasets,hf-datasets]
+# if IN_COLAB:
+#   !pip install pandas numpy matplotlib seaborn pillow scikit-learn tensorflow
+#   !pip install --upgrade kagglehub[pandas-datasets,hf-datasets]
 
 import kagglehub
 
@@ -356,7 +356,7 @@ x = Dense(512, activation='relu')(x) # Added another Dense layer
 x = Dense(256, activation='relu')(x) # Existing Dense layer
 predictions = Dense(1, activation='sigmoid')(x) # Output layer for binary classification
 
-model = Model(inputs=base_model.input, outputs=predictions)
+model = Model(inputs=inputs, outputs=predictions)
 model.compile(optimizer=Adam(learning_rate=0.0001), loss='binary_crossentropy', metrics=['accuracy'])
 
 """## 4. Data Generators