TabTransformer dimensions issue

[duncanmcelfresh]

this occurs on the following datasets:

• openml__balance-scale__11
• openml__ecoli__145977
• openml__jungle_chess_2pcs_raw_endgame_complete__167119
• openml__mfeat-morphological__18

traceback:

Traceback (most recent call last):
  File "/home/shared/tabzilla/TabSurvey/tabzilla_experiment.py", line 137, in __call__
    result = cross_validation(model, self.dataset, self.time_limit)
  File "/home/shared/tabzilla/TabSurvey/tabzilla_utils.py", line 236, in cross_validation
    loss_history, val_loss_history = curr_model.fit(
  File "/home/shared/tabzilla/TabSurvey/models/tabtransformer.py", line 108, in fit
    out = self.model(x_categ, x_cont)
  File "/opt/conda/envs/torch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/envs/torch/lib/python3.10/site-packages/torch/nn/parallel/data_parallel.py", line 166, in forward
    return self.module(*inputs[0], **kwargs[0])
  File "/opt/conda/envs/torch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/shared/tabzilla/TabSurvey/models/tabtransformer.py", line 580, in forward
    return self.mlp(x)
  File "/opt/conda/envs/torch/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/shared/tabzilla/TabSurvey/models/tabtransformer.py", line 469, in forward
    x = torch.softmax(x, dim=1)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)