Setting Loss Function in Tabular learner

Hey !
I am having difficulty in changing the loss function for my tabular learner. What I am trying to do is as follows :

  1. First train the model to predict regressionaly on a dataset.
  2. Then changing the data set such that now it becomes a classification problem and now using transfer learning to train the model trained earlier to classify on this dataset.

When I am manually trying to change the loss function by
learn.loss_func= CrossEntropyFlat()

And then running lr_find,

I am getting such an error :
cuda runtime error (710) : device-side assert triggered at /opt/conda/conda-bld/pytorch_1579022034529/work/torch/csrc/generic/serialization.cpp:148

Please help me in solving this problem !