Unable to do hyperparamters search on sentence-pair classification(regression) task

[Gregory5949]

Describe the bug
Hello, dear developers of simpletransformers!

Error 'wandb: ERROR Run xg0wdljm errored: '(wandb: ERROR Run 4i4jmyqn errored: ValueError('Target size (torch.Size([2])) must be the same as input size (torch.Size([2, 2]))'))' occurs when I'm trying script https://simpletransformers.ai/docs/tips-and-tricks/ ('6. Putting it all together') but for sentence-pair classification(regression). As I see, it's rather an issue of wandb, that's why I couldn't specify a class causing it.

To Reproduce

import logging

import pandas as pd
import sklearn

import wandb
from simpletransformers.classification import (
    ClassificationArgs,
    ClassificationModel,
)

sweep_config = {
    "method": "bayes",  # grid, random
    "metric": {"name": "train_loss", "goal": "minimize"},
    "parameters": {
        "num_train_epochs": {"values": [2, 3, 5]},
        "learning_rate": {"min": 5e-5, "max": 4e-4},
    },
}

sweep_id = wandb.sweep(sweep_config, project="Simple Sweep")

logging.basicConfig(level=logging.INFO)
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)


train_data = [
    [
        "Aragorn was the heir of Isildur",
        "Gimli fought with a battle axe",
        1,
    ],
    [
        "Frodo was the heir of Isildur",
        "Legolas was an expert archer",
        0,
    ],
]
train_df = pd.DataFrame(train_data)
train_df.columns = ["text_a", "text_b", "labels"]

# Preparing eval data
eval_data = [
    [
        "Theoden was the king of Rohan",
        "Gimli's preferred weapon was a battle axe",
        1,
    ],
    [
        "Merry was the king of Rohan",
        "Legolas was taller than Gimli",
        0,
    ],
]
eval_df = pd.DataFrame(eval_data)
eval_df.columns = ["text_a", "text_b", "labels"]


model_args = ClassificationArgs()

model_args.save_best_model = True
model_args.reprocess_input_data = True
model_args.regression = True
model_args.use_early_stopping = True
model_args.early_stopping_delta = 0.01
model_args.early_stopping_metric = "mae"
model_args.early_stopping_metric_minimize = True
model_args.early_stopping_patience = 5
model_args.evaluation_during_training = True
model_args.evaluate_during_training_steps = 100
model_args.overwrite_output_dir = True
model_args.evaluate_during_training = True
model_args.manual_seed = 42
model_args.wandb_project = "Simple Sweep"

def train():
    # Initialize a new wandb run
    wandb.init()

    # Create a TransformerModel
    model = ClassificationModel(
        "bert",
        "deeppavlov/rubert-base-cased",
        use_cuda=True,
        args=model_args,
        sweep_config=wandb.config,
    )

    # Train the model
    model.train_model(train_df, eval_df=eval_df)
    
    # Evaluate the model
    model.eval_model(eval_df=eval_df)

    

    # Sync wandb
    wandb.join()


wandb.agent(sweep_id, train)

Expected behavior
In my assumption it should work like in script https://simpletransformers.ai/docs/tips-and-tricks/ ('6. Putting it all together')

Screenshots

Desktop (please complete the following information):

• Google colab

Sincerely,
Grigory