Estimator should be an estimator implementing 'fit' method,

6

Change to criterion = torch.nn.CrossEntropyLoss,
Similar issue here.