Update evaluator.py

Author: codelion

examples/mlx_kernel_optimization/evaluator.py

@@ -306,16 +306,9 @@ def loss_fn(model, inputs, targets):
             logits_flat = logits.reshape(-1, vocab_size)
             targets_flat = targets.reshape(-1)
 
-            # Mask padding tokens (assume 0 is pad token)
-            mask = targets_flat != 0
-            if mx.sum(mask) == 0:  # All padding, use all tokens
-                mask = mx.ones_like(targets_flat, dtype=mx.bool_())
-            
-            # Apply mask
-            logits_masked = logits_flat[mask]
-            targets_masked = targets_flat[mask]
-            
-            return nn.losses.cross_entropy(logits_masked, targets_masked, reduction='mean')
+            # Simple cross-entropy without masking to avoid boolean indexing issues
+            # MLX doesn't support boolean indexing, so we'll compute loss on all tokens
+            return nn.losses.cross_entropy(logits_flat, targets_flat, reduction='mean')
 
         # Gradient function
         value_and_grad_fn = mx.value_and_grad(loss_fn)
@@ -324,9 +317,17 @@ def loss_fn(model, inputs, targets):
         def get_memory_usage():
             # Simple memory estimation based on array sizes
             total_memory = 0
-            for param in model.parameters():
-                if hasattr(param, 'size'):
-                    total_memory += param.size * 4  # Assume 4 bytes per float
+            try:
+                for param in model.parameters():
+                    if hasattr(param, 'shape'):
+                        # Calculate memory usage: shape -> total elements -> bytes
+                        total_elements = 1
+                        for dim in param.shape:
+                            total_elements *= dim
+                        total_memory += total_elements * 4  # Assume 4 bytes per float32
+            except Exception:
+                # Fallback to simple estimation
+                total_memory = 64 * 1024 * 1024  # 64MB default
             return total_memory / (1024 * 1024)  # MB
 
         initial_memory = get_memory_usage()