segmentaion fault occurs in convolution(group=in_channel) during mixed precision training

[HiromichiKamata]

When performing mixed precision training, I got segmentation fault when I specified the convolution group=in_channel.

import numpy as np
import nnabla as nn
import nnabla.functions as F  
import nnabla.parametric_functions as PF
from nnabla.ext_utils import get_extension_context


ctx = get_extension_context(
        "cudnn",
        device_id='0',
        type_config='half'
    )

nn.set_default_context(ctx)
x = nn.Variable((8, 32, 32, 32), need_grad=True)
x.d = np.random.random(x.shape)

h = PF.convolution(x, 32, (3,3), group=32)
loss = F.sum(h)
loss.forward(function_post_hook=lambda f: print(f'forward {f}'))
loss.backward(function_post_hook=lambda f: print(f'backward {f}'))

Above sample code outputs below:

2022-04-20 16:49:26,777 [nnabla][INFO]: Initializing CPU extension...
2022-04-20 16:49:26,945 [nnabla][INFO]: Initializing CUDA extension...
2022-04-20 16:49:26,971 [nnabla][INFO]: Initializing cuDNN extension...
forward ConvolutionCudaCudnn
forward SumCuda
backward SumCuda
Segmentation fault (core dumped)

For some insights, if I set channel_last option, like h = PF.convolution(x, 32, (3,3), group=32, channel_last=True), this segmentation fault does not occur.