Hi.
I have the system shutdown problem while training. After training a few batches, the computer shut down itself and rebooted. I think the custom layer causes this problem because training works well without that layer. The custom layer applies self-attention to compute the weighted mean and standard deviation in the time domain.
I trained with 4 GPUs using nn.DataParallel and the Pytorch version is 1.5.0.
class AttnPooling(nn.Module):
    def __init__(self, din, dh=500, eps=1e-12):
        super(AttnPooling, self).__init__()
        self.eps = eps
        self.w1 = Parameter(torch.Tensor(dh, din))
        self.w2 = Parameter(torch.Tensor(1, dh))
        self.reset_parameters()

    def reset_parameters(self) -> None:
        init.kaiming_uniform_(self.w1, a=math.sqrt(5))
        init.kaiming_uniform_(self.w2, a=math.sqrt(5))

    def forward(self, x, dim):
        # x : (Batch, F_dim, Time)
        # h : (Batch, Time, F_dim)
        h = x.transpose(1, 2)
        attn = F.relu(F.linear(h, self.w1))
        attn = F.softmax(F.linear(attn, self.w2), dim=1)

        mean = torch.bmm(x, attn)
        variance = (x - mean).pow(2).mean(dim=-1)
        mean = mean.squeeze(dim)

        mask = (variance <= self.eps).type(variance.dtype).to(variance.device)
        variance = (1.0 - mask) * variance + mask * self.eps
        stddev = variance.sqrt()

        # mean: (B, F_dim), stddev: (B, F_dim)
        pooling = torch.cat((mean, stddev), dim=-1)
        return pooling