Add few tests to main.py

Improve calculating activations
Scale down model and therefore compute and disk requirements
Reduce reserved memory for slurm task, as it was not needed