get errors when apply apex to DARTS

hi,

Since apex works well with ResNets, and DenseNets, I want to apply apex to train_search.py in https://github.com/MandyMo/DARTS.

But I get some errors.

If only use fp16, by inserting
self.model, [self.w_opt,self.alpha_opt ] = amp.initialize(self.model, [self.w_opt,self.alpha_opt ], 'O1')
after Line114 https://github.com/MandyMo/DARTS/blob/master/train_search.py#L114 and replacing Line433 by
with amp.scale_loss(train_loss, w_opt) as scaled_loss:
for the default setting (arg.mode=3), the error is

Traceback (most recent call last):
File "train01_search.py", line 458, in
trainer.train()
File "train01_search.py", line 451, in train
self.train_3()
File "train01_search.py", line 371, in train_3
train_logits = self.model(train_input)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 378, in forward
outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 399, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
raise output
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
output = module(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/model_search.py", line 92, in forward
s0, s1 = s1, cell(s0, s1, w_alpha_reduction)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/model_search.py", line 51, in forward
stats.append(sum(self._ops[i_n](stats[i_n-s_n], weights[i_n]) for i_n in range(s_n, e_n)))
File "/home/2019wea/code01/model_search.py", line 51, in
stats.append(sum(self._ops[i_n](stats[i_n-s_n], weights[i_n]) for i_n in range(s_n, e_n)))
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/model_search.py", line 25, in forward
return sum(w * op(x) for w, op in zip(weights, self._ops))
File "/home/2019wea/code01/model_search.py", line 25, in
return sum(w * op(x) for w, op in zip(weights, self._ops))
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/operators.py", line 116, in forward
out = self.bn(out)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/batchnorm.py", line 83, in forward
exponential_average_factor, self.eps)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/functional.py", line 1697, in batch_norm
training, momentum, eps, torch.backends.cudnn.enabled
RuntimeError: expected scalar type Half but found Float

If apply fp16 and syncbn to train_search.py, the error is

Traceback (most recent call last):
File "train01_search.py", line 458, in
trainer.train()
File "train01_search.py", line 451, in train
self.train_3()
File "train01_search.py", line 371, in train_3
train_logits = self.model(train_input)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 378, in forward
outputs = self.parallel_apply(self._module_copies[:len(inputs)], inputs, kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/distributed.py", line 399, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 83, in parallel_apply
raise output
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/parallel/parallel_apply.py", line 59, in _worker
output = module(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/model_search.py", line 94, in forward
s0, s1 = s1, cell(s0, s1, w_alpha_normal)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/model_search.py", line 47, in forward
stats = [self.pre0(s0), self.pre1(s1)]
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/home/2019wea/code01/operators.py", line 29, in forward
return self.op(x)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/container.py", line 92, in forward
input = module(input)
File "/opt/anaconda3/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in call
result = self.forward(*input, **kwargs)
File "/opt/anaconda3/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm.py", line 85, in forward
return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.train
ing or not self.track_running_stats, exponential_average_factor, self.process_group, self.channel_last, self.fuse_relu)
File "/opt/anaconda3/lib/python3.6/site-packages/apex/parallel/optimized_sync_batchnorm_kernel.py", line 47, in forward
r_m_inc = mean if running_mean.dtype != torch.float16 else mean.half()
AttributeError: 'NoneType' object has no attribute 'dtype'

In many classes, the outputs of forward are operators instead of tensors. I am not sure, but this seems to be the reason for errors.

Reproduce the two failures using train_search.py, starting around line 114 where amp.initialize is added and line 433 where amp.scale_loss is used. Read model_search.py and operators.py at the stack-trace locations, then compare the fp16 and fp16-with-syncbn paths. Done would require identifying and documenting a compatible path that completes training without either reported error.