add support for custom partition

Question from Sampo:

> Quick question regarding oellm-cli: is it possible to specify the partition to use, and if not, could this be added as a feature? I'm doing a series of relatively quick tests on LUMI at the moment and small-g is crowded, so it would be nice to be able to use dev-g.

We should just add a flag that allows to override the default partition used in the cluster.