diff --git a/EIANN/rules/dendritic_loss.py b/EIANN/rules/dendritic_loss.py index 54b7892a..5b9a9f5b 100644 --- a/EIANN/rules/dendritic_loss.py +++ b/EIANN/rules/dendritic_loss.py @@ -5,13 +5,7 @@ class DendriticLossBias(BiasLearningRule): def step(self): # self.population.bias.data += self.learning_rate * self.population.dendritic_state - bias_update += self.learning_rate * self.population.plateau - - # Convert to same dtype as bias for AMP compatibility - if self.population.network.use_amp and bias_update.dtype != self.population.bias.dtype: - bias_update = bias_update.to(self.population.bias.dtype) - - self.population.bias.data += bias_update + self.population.bias.data += self.learning_rate * self.population.plateau class DendriticLoss_6(LearningRule): @@ -31,9 +25,5 @@ def step(self): delta_weight = torch.outer( torch.clamp(self.projection.post.forward_dendritic_state.detach().clone(), min=-1, max=1), torch.clamp(self.projection.pre.forward_prev_activity, min=0, max=1)) - - # Ensure delta_weight has same dtype as weight for AMP compatibility - if self.projection.post.network.use_amp and delta_weight.dtype != self.projection.weight.dtype: - delta_weight = delta_weight.to(self.projection.weight.dtype) self.projection.weight.data += self.sign * self.learning_rate * delta_weight diff --git a/EIANN/simulate/jobscripts/run_EIANN_gpu_amarel_mnist_ray.sh b/EIANN/simulate/jobscripts/run_EIANN_gpu_amarel_mnist_ray.sh new file mode 100644 index 00000000..5469bb68 --- /dev/null +++ b/EIANN/simulate/jobscripts/run_EIANN_gpu_amarel_mnist_ray.sh @@ -0,0 +1,69 @@ +#!/bin/bash + +# Export the current date and time for job labeling +export DATE=$(date +%Y%m%d_%H%M%S) +export LABEL="$1" +export JOB_NAME=eiann_gpu_mnist_ray_"$LABEL"_"$DATE" + +# Environment variables to optimize performance +export OMP_NUM_THREADS=4 +export MKL_NUM_THREADS=1 +export NUMEXPR_NUM_THREADS=1 +export OPENBLAS_NUM_THREADS=1 +export CUDA_DEVICE_MAX_CONNECTIONS=1 +export WANDB_START_METHOD=thread + +# Create directories for logs and scratch data +mkdir -p /scratch/${USER}/logs/eiann +mkdir -p /scratch/${USER}/data/eiann + +# Submit the job +sbatch <