Thank you for your reply, and sorry for my late reply.



 ebarsoum:

It looks like it fails in tracing,


Yes, torch.jit.trace doesn’t work.
>>> torch.jit.trace(model,dummy_input)
/home/XXXX/darts/cnn/eval-EXP-20200710-150423/utils.py:105: TracerWarning: Converting a tensor to a Python index might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  mask = Variable(torch.cuda.FloatTensor(x.size(0), 1, 1, 1).bernoulli_(keep_prob))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 875, in trace
    check_tolerance, _force_outplace, _module_class)
  File "/home/XXXX/darts/cnn/eval-EXP-20200710-150423/venv/lib/python3.6/site-packages/torch/jit/__init__.py", line 1027, in trace_module
    module._c._create_method_from_trace(method_name, func, example_inputs, var_lookup_fn, _force_outplace)
RuntimeError: Only tensors, lists and tuples of tensors can be output from traced functions

If I remove logits_aux following ptrblck’s advice, it work well.
I don’t know why. So, I would appreciate it if you could inform when you find it.
Thank you!