Allow definition of task heads other than MLP/CNN for ESM2 finetuning

Allow definition of task heads other than MLP/CNN for ESM2 finetuning #3604