update code proc

Author: kchu25

src/final_and_code/_helpers.jl

@@ -61,7 +61,7 @@ function _processor_train_step!(processor, opt_state, model, seq, proc_wrap, inf
     preds, grad = _compute_code_gradients(model, code, proc_wrap, inf_layer, predict_position)
 
     loss, grads = Flux.withgradient(processor) do p
-        proc_grad = proc_wrap.process_code(p, code, grad; step=step)
+        proc_grad = p(code, grad; step=step)
         sum(abs2, vec(sum(proc_grad .* code, dims=(1,2))) - preds)
     end
 

src/final_and_code/code_processor_eval.jl

@@ -100,8 +100,10 @@ function evaluate_processor(model, processor, dataloader, proc_wrap, set_name::S
         if isnothing(gyro_shape)
             gyro_shape = size(gyro)
         end
-        
-        proc_gyro = proc_wrap.process_code(processor, code, gyro; training=false)
+
+        processor.training[] = false  # Set processor to eval mode
+
+        proc_gyro = processor(code, gyro)
 
         _accumulate_batch_stats!(stats, code, gyro, proc_gyro, preds, epsilon)
     end