more bug fixes

app.py

@@ -24,7 +24,7 @@ for preset in model_presets:
     chat_template = get_appropriate_chat_template(preset)
     chat_state = ChatState(model, "", chat_template)
     prompt, response = chat_state.send_message("Hello")
-    print("model " + preset + "loaded and initialized.")
+    print("model " + preset + " loaded and initialized.")
     print("The model responded: " + response)
     models.append(model)
 

models.py

@@ -41,10 +41,8 @@ def get_default_layout_map(preset_name, device_mesh):
 def log_applied_layout_map(model):
     if "Gemma" in type(model).__name__:
         transformer_decoder_block_name = "decoder_block_1"
-    elif "Llama3" in type(model).__name__ or "Mistral" in type(model).__name__:
+    else:  # works for Llama, Mistral, Vicuna
         transformer_decoder_block_name = "transformer_layer_1"
-    else:
-        assert (0, "Model type not recognized. Cannot display model layout.")
 
     # See how layer sharding was applied
     embedding_layer = model.backbone.get_layer("token_embedding")