To address your points:

Can you paste the error with tensorboard and also which line you comment to solve the issue? Have you checked that this PR does not brake our existing build? Maybe add a parameter --ray_tune so that if we don't specify it, the behavior is the same as currently (I'm a bit skeptical regarding the new parameters).
Yes, here's the tensorboard error:
`
Serializing 'writer' <torch.utils.tensorboard.writer.SummaryWriter object at 0x7f1c93b6efd0>...
!!! FAIL serialization: cannot pickle 'tensorflow.python.lib.io._pywrap_file_io.WritableFile' object
Serializing '_annotated' FunctionTrainable...
================================================================================
Variable:

    FailTuple(writer [obj=<torch.utils.tensorboard.writer.SummaryWriter object at 0x7f1c93b6efd0>, parent=<function train_multitask at 0x7f1c93b644c0>])

was found to be non-serializable. There may be multiple other undetected variables that were non-serializable.
Consider either removing the instantiation/imports of these variables or moving the instantiation into the scope of the function/class.
`

Once I commented out all references to writer (eg. writer.add_scalar) in train_multitask and model_eval_multitask (which is called by train_multitask), the code runs fine.

Also could you give us some instructions on how to install raytune (It seems like you struggled, so I would love to get your opinion on how to do this :) )

Raytune was actually simple enough to install on AWS with pip. But here's the command:
pip install -U "ray[tune]" # installs Ray + dependencies for Ray Tune

Hi @marie-huynh raytune serves as a wrapper function and requires the training function (train_multitask) to be modified to accept a config dictionary passed by ray and use the arguments within it. So I don't think I can run raytune without directly editing multitask_classifier.

However, I have not made functional edits to train_multitask (only editing the input argument to be a config dictionary passed by ray instead of args). Internally, I've also edited train_multitask to use arguments from this config dictionary.

Finally, I've added two new command line flags:

• --tune_hyperparameters: if this isn't provided, training runs normally
• --num_tuning_runs: Sets the number of hyperparameter training experiments to run

I've updated my script to avoid out of memory errors by saving directly to the disk rather than RAM. As long as you update your disk size, it should run ok.