handle transformer dim mismatch during tuning

            raise optuna.TrialPruned(f"OOM error: {e}")

        else:

            raise e

    except AssertionError as e:

        # Cache dimension and n_head mismatch

        if "embed_dim must be divisible by num_heads" in str(e):

            logger.error("embed_dim must be divisible by num_heads")

            raise optuna.TrialPruned(f"embed_dim must be divisible by num_heads: {e}")

        else:

            raise e

    except Exception as e:

        loss = 1e6

        raise e