diff --git a/train_v2.py b/train_v2.py
index a4c3d67..57e82f8 100644
--- a/train_v2.py
+++ b/train_v2.py
@@ -23,6 +23,7 @@ configure_xla_paths()
 
 import tensorflow as tf
 import numpy as np
+print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))
 
 # ==========================================
 # 1. MODEL ARCHITECTURE
@@ -87,7 +88,12 @@ def produce_sample(model, vectorize_layer, vocab, seed, num_generate=300, temper
     # 1. Safely tokenize the seed character-by-character using the vectorizer
     # We split the string into individual characters so the vectorizer processes them properly
     seed_chars = tf.strings.unicode_split(seed, input_encoding='UTF-8')
-    input_ids = vectorize_layer(seed_chars).numpy().tolist()
+
+    input_ids = tf.squeeze(vectorize_layer(seed_chars)) 
+    if input_ids.ndim == 0:  # Safeguard if the seed is a single character
+        input_ids = tf.expand_dims(input_ids, 0)
+        
+    #input_ids = vectorize_layer(seed_chars).numpy().tolist()
 
     generated_ids = []
     states = None
@@ -151,11 +157,12 @@ def main():
     DEFAULT_NUM_LAYERS = 3
     DEFAULT_DROPOUT_RATE = 0.4
 
-    TEXT_FILE = "emily_post.txt"
+    #TEXT_FILE = "emily_post.txt"
+    TEXT_FILE = "combined_training_data.txt"
     VOCAB_FILE = "vocab.txt"
     SEQ_LENGTH = 512
     BATCH_SIZE = 128
-    EPOCHS_PER_RUN = 20
+    EPOCHS_PER_RUN = 22
     BUFFER_SIZE = 10000
     SEED_TEXT = "In the case of unlawful enterance"
 
@@ -208,11 +215,16 @@ def main():
 
     # Initialize model dynamically using the determined metrics
     model = CharacterTextModel(vocab_size=vocab_size, embedding_dim=emb_dim, rnn_units=rnn_units, num_layers=num_layers, dropout_rate=dropout_rate)
-    model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH))
+
+    dummy_input = tf.zeros((BATCH_SIZE, SEQ_LENGTH), dtype=tf.int32)
+    _ = model(dummy_input, training=False)
+    
+    #model.build(input_shape=(BATCH_SIZE, SEQ_LENGTH))
 
     loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
     model.compile(optimizer='adam', loss=loss_fn)
 
+
     if is_resuming:
         print(f"[Weights] Restoring model layer states from: {args.modelfile}")
         model.load_weights(args.modelfile)
@@ -271,7 +283,8 @@ def main():
         total_epochs = initial_epoch + EPOCHS_PER_RUN
         print(f"[Train] Commencing training loop inside checkpoint tree: {checkpoint_dir}")
         print(f"[Train] Running Epochs {initial_epoch + 1} through {total_epochs}...")
-
+        print(f"[DEBUG] Verification -> initial_epoch: {initial_epoch}, total_epochs: {total_epochs}")
+        
         model.fit(
             dataset,
             epochs=total_epochs,