Repository: autogluon/autogluon
Issue: [BUG] Exception during fitting in calibrate_model() when using custom metric with needs_proba=True

- [ ] I have checked that this bug exists on the latest stable version of AutoGluon
- [ ] and/or I have checked that this bug exists on the latest mainline of AutoGluon via source installation

**Describe the bug**
At the end of training an exception is raised during `calibrate_model()` when using a custom metric with `needs_proba=True`

I'm using version 0.6.0 (I have a different reported issue with 0.6.1).

**Expected behavior**
No exception

**To Reproduce**
```python
import numpy as np

from autogluon.core.metrics import make_scorer
from autogluon.tabular import TabularDataset, TabularPredictor

def custom_score(y_true, y_proba):
    pred_true_np  = y_true[y_proba >= 0.6]
    pred_false_np = y_true[y_proba <= 0.4]

    pred_true_correct_count  =                      np.count_nonzero(pred_true_np)
    pred_false_correct_count = pred_false_np.size - np.count_nonzero(pred_false_np)

    pred_count         = pred_true_np.size       + pred_false_np.size
    pred_correct_count = pred_true_correct_count + pred_false_correct_count

    if pred_count == 0:
        return 0.5

    score = pred_correct_count / pred_count
    return score

metric = make_scorer(
    name="custom_score",
    score_func=custom_score,
    optimum=1,
    greater_is_better=True,
    needs_proba=True,
)

train_data = TabularDataset('https://autogluon.s3.amazonaws.com/datasets/Inc/train.csv')
test_data = TabularDataset('https://autogluon.s3.amazonaws.com/datasets/Inc/test.csv')
predictor = TabularPredictor(label='class', eval_metric=metric).fit(train_data=train_data, time_limit=60)
```

**Logs**
```python
Beginning AutoGluon training ... Time limit = 60s
AutoGluon Version:  0.6.0
Python Version:     3.9.13
Operating System:   Windows
Platform Machine:   AMD64
Platform Version:   10.0.19045
Train Data Rows:    39073
Train Data Columns: 14
Label Column: class
Preprocessing data ...
AutoGluon infers your prediction problem is: 'binary' (because only two unique label-values observed).
        2 unique label values:  [' <=50K', ' >50K']
        If 'binary' is not the correct problem_type, please manually specify the problem_type parameter during predictor init (You may specify problem_type as one of: ['binary', 'multiclass', 'regression'])
Selected class <--> label mapping:  class 1 =  >50K, class 0 =  <=50K
        Note: For your binary classification, AutoGluon arbitrarily selected which label-value represents positive ( >50K) vs negative ( <=50K) class.
        To explicitly set the positive_class, either rename classes to 1 and 0, or specify positive_class in Predictor init.
Using Feature Generators to preprocess the data ...
Fitting AutoMLPipelineFeatureGenerator...
        Inferring data type of each feature based on column values. Set feature_metadata_in to manually specify special dtypes of the features.
        Stage 1 Generators:
                Fitting AsTypeFeatureGenerator...
                        Note: Converting 1 features to boolean dtype as they only contain 2 unique values.
        Stage 2 Generators:
                Fitting FillNaFeatureGenerator...
        Stage 3 Generators:
                Fitting IdentityFeatureGenerator...
                Fitting CategoryFeatureGenerator...
                        Fitting CategoryMemoryMinimizeFeatureGenerator...
        Stage 4 Generators:
                Fitting DropUniqueFeatureGenerator...
        Types of features in original data (raw dtype, special dtypes):
                ('int', [])    : 6 | ['age', 'fnlwgt', 'education-num', 'capital-gain', 'capital-loss', ...]
                ('object', []) : 8 | ['workclass', 'education', 'marital-status', 'occupation', 'relationship', ...]
        Types of features in processed data (raw dtype, special dtypes):
                ('category', [])  : 7 | ['workclass', 'education', 'marital-status', 'occupation', 'relationship', ...]
                ('int', [])       : 6 | ['age', 'fnlwgt', 'education-num', 'capital-gain', 'capital-loss', ...]
                ('int', ['bool']) : 1 | ['sex']
        0.4s = Fit runtime
        14 features in original data used to generate 14 features in processed data.
Data preprocessing and feature engineering runtime = 0.45s ...
AutoGluon will gauge predictive performance using evaluation metric: 'custom_score'
        This metric expects predicted probabilities rather than predicted class labels, so you'll need to use predict_proba() instead of predict()
        To change this, specify the eval_metric parameter of Predictor()
Automatically generating train/validation split with holdout_frac=0.0639828014229775, Train Rows: 36573, Val Rows: 2500
Fitting 13 L1 models ...
Fitting model: KNeighborsUnif ... Training model for up to 59.55s of the 59.55s of remaining time.
        0.7752   = Validation score   (custom_score)
        1.66s    = Training   runtime
        0.01s    = Validation runtime
Fitting model: KNeighborsDist ... Training model for up to 57.87s of the 57.87s of remaining time.
        0.7932   = Validation score   (custom_score)
        0.05s    = Training   runtime
        0.02s    = Validation runtime
Fitting model: LightGBMXT ... Training model for up to 57.79s of the 57.78s of remaining time.
        0.9151   = Validation score   (custom_score)
        3.95s    = Training   runtime
        0.05s    = Validation runtime
Fitting model: LightGBM ... Training model for up to 53.71s of the 53.7s of remaining time.
        0.919    = Validation score   (custom_score)
        1.21s    = Training   runtime
        0.02s    = Validation runtime
Fitting model: RandomForestGini ... Training model for up to 52.46s of the 52.46s of remaining time.
        0.8937   = Validation score   (custom_score)
        3.86s    = Training   runtime
        0.18s    = Validation runtime
Fitting model: RandomForestEntr ... Training model for up to 47.89s of the 47.89s of remaining time.
        0.8966   = Validation score   (custom_score)
        4.65s    = Training   runtime
        0.18s    = Validation runtime
Fitting model: CatBoost ... Training model for up to 42.59s of the 42.58s of remaining time.
        Ran out of time, early stopping on iteration 487.
        0.9161   = Validation score   (custom_score)
        42.64s   = Training   runtime
        0.02s    = Validation runtime
Fitting model: WeightedEnsemble_L2 ... Training model for up to 59.55s of the -0.13s of remaining time.
        0.9202   = Validation score   (custom_score)
        0.34s    = Training   runtime
        0.0s     = Validation runtime
AutoGluon training complete, total runtime = 60.55s ... Best model: "WeightedEnsemble_L2"
Traceback (most recent call last):
  File "test.py", line 32, in <module>
    predictor = TabularPredictor(label='class', eval_metric=metric).fit(train_data=train_data, time_limit=60)
  File ".venv\\lib\\site-packages\\autogluon\\core\\utils\\decorators.py", line 30, in _call
    return f(*gargs, **gkwargs)
  File ".venv\\lib\\site-packages\\autogluon\\tabular\\predictor\\predictor.py", line 868, in fit
    self._post_fit(
  File ".venv\\lib\\site-packages\\autogluon\\tabular\\predictor\\predictor.py", line 921, in _post_fit
    self._trainer.calibrate_model()
  File ".venv\\lib\\site-packages\\autogluon\\core\\trainer\\abstract_trainer.py", line 2953, in calibrate_model
    temp_scalar = tune_temperature_scaling(y_val_probs=y_val_probs, y_val=y_val,
  File ".venv\\lib\\site-packages\\autogluon\\core\\calibrate\\temperature_scaling.py", line 57, in tune_temperature_scaling
    optimizer.step(temperature_scale_step)
  File ".venv\\lib\\site-packages\\torch\\optim\\lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\optim\\optimizer.py", line 113, in wrapper
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\autograd\\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\optim\\lbfgs.py", line 311, in step
    orig_loss = closure()
  File ".venv\\lib\\site-packages\\torch\\autograd\\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\autogluon\\core\\calibrate\\temperature_scaling.py", line 52, in temperature_scale_step
    loss = nll_criterion(new_logits, y_val_tensor)
  File ".venv\\lib\\site-packages\\torch\\nn\\modules\\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\nn\\modules\\loss.py", line 1164, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File ".venv\\lib\\site-packages\\torch\\nn\\functional.py", line 3014, in cross_entropy
    return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
RuntimeError: expected scalar type Long but found Int
```

**Installed Versions**
autogluon-0.6.0

**Additional context**
In the actual training I'm doing I get a slightly different exception than above:

```python
Traceback (most recent call last):
  File "ag.py", line 79, in train
    predictor.fit(train_data)
  File ".venv\\lib\\site-packages\\autogluon\\core\\utils\\decorators.py", line 30, in _call
    return f(*gargs, **gkwargs)
  File ".venv\\lib\\site-packages\\autogluon\\tabular\\predictor\\predictor.py", line 868, in fit
    self._post_fit(
  File ".venv\\lib\\site-packages\\autogluon\\tabular\\predictor\\predictor.py", line 921, in _post_fit
    self._trainer.calibrate_model()
  File ".venv\\lib\\site-packages\\autogluon\\core\\trainer\\abstract_trainer.py", line 2953, in calibrate_model
    temp_scalar = tune_temperature_scaling(y_val_probs=y_val_probs, y_val=y_val,
  File ".venv\\lib\\site-packages\\autogluon\\core\\calibrate\\temperature_scaling.py", line 57, in tune_temperature_scaling
    optimizer.step(temperature_scale_step)
  File ".venv\\lib\\site-packages\\torch\\optim\\lr_scheduler.py", line 65, in wrapper
    return wrapped(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\optim\\optimizer.py", line 113, in wrapper
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\autograd\\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\torch\\optim\\lbfgs.py", line 311, in step
    orig_loss = closure()
  File ".venv\\lib\\site-packages\\torch\\autograd\\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File ".venv\\lib\\site-packages\\autogluon\\core\\calibrate\\temperature_scaling.py", line 50, in temperature_scale_step
    temp = temperature_param.unsqueeze(1).expand(logits.size(0), logits.size(1))
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
```

Changed files:
core/src/autogluon/core/trainer/abstract_trainer.py, tabular/tests/unittests/edgecases/test_edgecases.py

diff --git a/core/src/autogluon/core/trainer/abstract_trainer.py b/core/src/autogluon/core/trainer/abstract_trainer.py
index 4e8c88693..4db69b800 100644
--- a/core/src/autogluon/core/trainer/abstract_trainer.py
+++ b/core/src/autogluon/core/trainer/abstract_trainer.py
@@ -95,6 +95,7 @@ class AbstractTrainer:
 
         self._num_rows_train = None
         self._num_cols_train = None
+        self._num_rows_val = None
 
         self.is_data_saved = False
         self._X_saved = False
@@ -128,6 +129,11 @@ class AbstractTrainer:
     def path_data(self) -> str:
         return self.path_utils + 'data' + os.path.sep
 
+    @property
+    def has_val(self) -> bool:
+        """Whether the trainer uses validation data"""
+        return self._num_rows_val is not None
+
     def load_X(self):
         if self._X_saved:
             path = self.path_data + 'X.pkl'
@@ -1968,7 +1974,7 @@ class AbstractTrainer:
             self._groups = groups
         self._num_rows_train = len(X)
         if X_val is not None:
-            self._num_rows_train += len(X_val)
+            self._num_rows_val = len(X_val)
         self._num_cols_train = len(list(X.columns))
         model_names_fit = self.train_multi_levels(X, y, hyperparameters=hyperparameters, X_val=X_val, y_val=y_val,
                                                   X_unlabeled=X_unlabeled, level_start=1, level_end=num_stack_levels+1, time_limit=time_limit, **kwargs)
@@ -2455,6 +2461,7 @@ class AbstractTrainer:
         time_train_start = self._time_train_start
         num_rows_train = self._num_rows_train
         num_cols_train = self._num_cols_train
+        num_rows_val = self._num_rows_val
         num_classes = self.num_classes
         # TODO:
         #  Disk size of models
@@ -2472,6 +2479,7 @@ class AbstractTrainer:
             'time_train_start': time_train_start,
             'num_rows_train': num_rows_train,
             'num_cols_train': num_cols_train,
+            'num_rows_val': num_rows_val,
             'num_classes': num_classes,
             'problem_type': problem_type,
             'eval_metric': eval_metric,
@@ -2939,13 +2947,13 @@ class AbstractTrainer:
             if m_full == model_name:
                 model_name_og = m
                 break
-        if self.bagged_mode:
-            y_val_probs = self.get_model_oof(model_name_og)
-            y_val = self.load_y().to_numpy()
-        else:
+        if self.has_val:
             X_val = self.load_X_val()
             y_val_probs = self.predict_proba(X_val, model_name_og)
             y_val = self.load_y_val().to_numpy()
+        else:  # bagged mode
+            y_val_probs = self.get_model_oof(model_name_og)
+            y_val = self.load_y().to_numpy()
 
         if self.problem_type == BINARY:
             # Convert one-dimensional array to be in the form of a 2-class multiclass predict_proba output
diff --git a/tabular/tests/unittests/edgecases/test_edgecases.py b/tabular/tests/unittests/edgecases/test_edgecases.py
index 640252459..03702225c 100644
--- a/tabular/tests/unittests/edgecases/test_edgecases.py
+++ b/tabular/tests/unittests/edgecases/test_edgecases.py
@@ -93,3 +93,29 @@ def test_num_folds_hpo(fit_helper):
     assert leaderboard.iloc[0]['num_models'] == 4
     assert leaderboard.iloc[1]['num_models'] == 4
     shutil.rmtree(predictor.path, ignore_errors=True)
+
+
+def test_use_bag_holdout_calibrate(fit_helper):
+    """
+    Test that use_bag_holdout=True works for calibration
+    Ensures the bug is fixed in https://github.com/autogluon/autogluon/issues/2674
+    """
+    init_args = dict(
+        eval_metric='log_loss'
+    )
+
+    fit_args = dict(
+        hyperparameters={'DUMMY': {}},
+        num_bag_folds=2,
+        use_bag_holdout=True,
+        calibrate=True,
+    )
+
+    dataset_name = 'adult'
+    fit_helper.fit_and_validate_dataset(
+        dataset_name=dataset_name,
+        init_args=init_args,
+        fit_args=fit_args,
+        expected_model_count=2,
+        refit_full=False,
+    )