Top_k_accuracy error

I was trying to use top_k_accuracy in a learner, and got the following error:

RuntimeError: expand(torch.CharTensor{[2, 1, 1]}, size=[2, 1]): the number of sizes provided (2) must be greater or equal to the number of dimensions in the tensor (3)

It is hard to copy all of my code, but here is a stripped-down example that replicates it:

from functools import partial
top_1_accuracy = partial(top_k_accuracy, k=1)

from fastai.tabular.all import *

df = pd.DataFrame({'x': [1, 2, 3, 4], 'y': [2, 4, 6, 8], 'z': ['foo', 'bar', 'bar', 'bar']})

splits = RandomSplitter(valid_pct=.5, seed=5)(df)
pan = TabularPandas(df, cont_names=['x', 'y'], y_names='z', splits=splits)
dls = pan.dataloaders(bs=2)
learn = tabular_learner(dls, metrics=top_1_accuracy)

learn.fit_one_cycle(2, 1e-3)

Debugging through that method, it looks like after the unsqueeze operation, the tensor is the wrong size. If I re-define the top_k_accuracy method and just take out the unsqueeze, it works:

def top_1_accuracy(inp, targ, k=1, axis=-1):
    "Computes the Top-k accuracy (`targ` is in the top `k` predictions of `inp`)"
    inp = inp.topk(k=k, dim=axis)[1]
    targ = targ.expand_as(inp)
    return (inp == targ).sum(dim=-1).float().mean()

I’m not 100% sure what I’m doing wrong or why the unsqueeze is in there - is this a bug, or is it an issue with the way I’m setting up my learner?