RuntimeError: expected type torch.cuda.FloatTensor but got torch.cuda.HalfTensor

The traceback points to the notebook's fit() call and the BCEWithLogitsLoss invocation in the model's forward method. Start by reproducing that call with FP16_Optimizer enabled and inspect the logits and labels at the loss boundary. Done means the source of the tensor-type mismatch is isolated and a verified resolution is identified.

Hi,
I get this error when I try to compute a BCEWithLogitsLoss for my model. I use the FP16_Optimizer to compute the gradients so I don't understand where the problem comes from. Here is the full Traceback:

RuntimeError Traceback (most recent call last)
in ()
----> 1 fit(1)

in fit(num_epocs)
10 batch = tuple(t.to(device) for t in batch)
11 input_ids, input_mask, segment_ids, label_ids = batch
---> 12 loss = model(input_ids, segment_ids, input_mask, label_ids)
13 if n_gpu > 1:
14 loss = loss.mean() # mean() to average on multi-gpu.

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
487 result = self._slow_forward(*input, **kwargs)
488 else:
--> 489 result = self.forward(*input, **kwargs)
490 for hook in self._forward_hooks.values():
491 hook_result = hook(self, input, result)

in forward(self, input_ids, token_type_ids, attention_mask, labels)
52 if labels is not None:
53 loss_fct = BCEWithLogitsLoss()
---> 54 loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1, self.num_labels))
55 return loss
56 else:

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/module.py in call(self, *input, **kwargs)
487 result = self._slow_forward(*input, **kwargs)
488 else:
--> 489 result = self.forward(*input, **kwargs)
490 for hook in self._forward_hooks.values():
491 hook_result = hook(self, input, result)

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/modules/loss.py in forward(self, input, target)
593 self.weight,
594 pos_weight=self.pos_weight,
--> 595 reduction=self.reduction)
596
597

~/anaconda3/envs/fastai/lib/python3.7/site-packages/torch/nn/functional.py in binary_cross_entropy_with_logits(input, target, weight, size_average, reduce, reduction, pos_weight)
2075 raise ValueError("Target size ({}) must be the same as input size ({})".format(target.size(), input.size()))
2076
-> 2077 return torch.binary_cross_entropy_with_logits(input, target, weight, pos_weight, reduction_enum)
2078
2079

RuntimeError: expected type torch.cuda.FloatTensor but got torch.cuda.HalfTensor

apex doesn't return any error during any of its calls (FusedAdam, FP16_Optimizer). Any idea?
Thanks