Upload model.py

model.py

@@ -168,10 +168,14 @@ def main():
 
     print(f"Created dynamic PyTorch model with sections: {list(model.sections.keys())}")
 
+    # Print the model architecture
+    print(model)
+
     first_section = next(iter(model.sections.keys()))
     first_layer = model.sections[first_section][0]
     input_features = first_layer[0].in_features
 
+    # Ensure the input tensor size matches the expected input size
     sample_input = torch.randn(1, input_features)
     output = model(sample_input)
     print(f"Sample output shape: {output.shape}")