Update constants.ts

constants.ts

@@ -412,11 +412,13 @@ export const LAYER_DEFINITIONS: Record<LayerType, LayerDefinition> = {
   [LayerType.CUSTOM]: {
     type: LayerType.CUSTOM,
     label: 'Custom / Code',
-    description: 'Any PyTorch/Python Module',
+    description: 'Define or instantiate any PyTorch Module',
     category: 'Utility',
     parameters: [
-        { name: 'class_name', type: 'string', label: 'Class Name', default: 'MyCustomLayer', description: "e.g., 'torch.nn.Identity' or 'MyBlock'" },
-        { name: 'args', type: 'string', label: 'Arguments', default: '', description: "e.g., 'dim=128, p=0.1'" }
+        { name: 'class_name', type: 'string', label: 'Class Name', default: 'MyCustomLayer', description: "Name of the class to instantiate" },
+        { name: 'args', type: 'string', label: 'Arguments', default: '', description: "Constructor args (e.g. dim=128)" },
+        { name: 'imports', type: 'string', label: 'Imports', default: '', description: "Required imports (e.g. import torch.nn.functional as F)" },
+        { name: 'definition_code', type: 'text', label: 'Python Code (Def)', default: '', description: "Full class definition if strictly custom." }
     ]
   },
   [LayerType.IDENTITY]: {