issues with pr #368

PR #368 has some issues, when testing test_distilling_emo(), it has the following error:

`Traceback (most recent call last):
File "test_autohf_distilling.py", line 180, in
test_distilling_emo()
File "test_autohf_distilling.py", line 157, in test_distilling_emo
automl.retrain_from_log(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/automl.py", line 1416, in retrain_from_log
self._trained_estimator = self._state._train_with_config(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/automl.py", line 363, in _train_with_config
estimator, train_time = train_estimator(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/ml.py", line 613, in train_estimator
train_time = estimator.fit(X_train, y_train, budget, **fit_kwargs)
File "/data/xliu127/projects/hyperopt/FLAML/flaml/model.py", line 653, in fit
self._model.train()
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/trainer.py", line 1325, in train
tr_loss_step = self.training_step(model, inputs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/trainer.py", line 1884, in training_step
loss = self.compute_loss(model, inputs)
File "/data/xliu127/projects/hyperopt/FLAML/flaml/nlp/huggingface/trainer.py", line 23, in compute_loss
outputs = model(**inputs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/models/distilbert/modeling_distilbert.py", line 766, in forward
loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/loss.py", line 1150, in forward
return F.cross_entropy(input, target, weight=self.weight,
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/functional.py", line 2846, in cross_entropy
return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
IndexError: Target 5 is out of bounds.`

When testing test_distilling_hf() with gpu_per_trial set to 1, it has the following error:

`Traceback (most recent call last):
File "test_autohf_distilling.py", line 180, in
test_distilling_emo()
File "test_autohf_distilling.py", line 154, in test_distilling_emo
automl.fit(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/automl.py", line 2147, in fit
self._search()
File "/data/xliu127/projects/hyperopt/FLAML/flaml/automl.py", line 2751, in _search
) = self._state._train_with_config(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/automl.py", line 363, in _train_with_config
estimator, train_time = train_estimator(
File "/data/xliu127/projects/hyperopt/FLAML/flaml/ml.py", line 613, in train_estimator
train_time = estimator.fit(X_train, y_train, budget, **fit_kwargs)
File "/data/xliu127/projects/hyperopt/FLAML/flaml/model.py", line 653, in fit
self._model.train()
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/trainer.py", line 1325, in train
tr_loss_step = self.training_step(model, inputs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/trainer.py", line 1884, in training_step
loss = self.compute_loss(model, inputs)
File "/data/xliu127/projects/hyperopt/FLAML/flaml/nlp/huggingface/trainer.py", line 38, in compute_loss
output_tea = self.teacher(**inputs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 1540, in forward
outputs = self.bert(
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 992, in forward
embedding_output = self.embeddings(
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/transformers/models/bert/modeling_bert.py", line 214, in forward
inputs_embeds = self.word_embeddings(input_ids)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
return forward_call(*input, **kwargs)
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/modules/sparse.py", line 158, in forward
return F.embedding(
File "/data/installation/anaconda3/envs/tmp/lib/python3.8/site-packages/torch/nn/functional.py", line 2044, in embedding
return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper__index_select)`