tom:

includes


I try to use a cudnn backhand ctc loss, according to the doc，I write the follow script
# -*- coding: utf-8 -*-
import torch
from torch import nn

class Net(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Net, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(3, 3), padding=(1, 1)),
            nn.ReLU(),
            nn.BatchNorm2d(out_channels),
            nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
        )

    def forward(self, x):
        x = self.features(x)
        x = x.squeeze(dim=2)
        x = x.permute(2, 0, 1)
        return x


if __name__ == '__main__':
    device = torch.device('cuda:0')
    a = torch.zeros((2, 3, 3, 160)).to(device)
    net = Net(3, 10)
    net.to(device)
    log_probs = net(a).log_softmax(2)
    target_lengths = [30, 20]
    targets = torch.randint(1, 10, (sum(target_lengths),), dtype=torch.int)
    target_lengths = torch.Tensor(target_lengths).to(torch.int)
    input_lengths = torch.Tensor([80, 80]).to(torch.int)
    ctc = torch.nn.CTCLoss()
    res = ctc(log_probs, targets, input_lengths, target_lengths)
    print('log_probs: {}, type: {}'.format(log_probs.size(), log_probs.dtype))
    print('input_lengths: {}, type: {}'.format(input_lengths.size(), input_lengths.dtype))
    print('targets: {}, type: {}'.format(targets.size(), targets.dtype))
    print('target_lengths: {}, type: {}'.format(target_lengths.size(), target_lengths.dtype))
    print('res.grad_fn:', res.grad_fn)
    print('res.grad_fn:', res.grad_fn.name())

the output is
D:\Anaconda3\python.exe E:/work/flask_learn/test.py
log_probs: torch.Size([80, 2, 10]), type: torch.float32
input_lengths: torch.Size([2]), type: torch.int32
targets: torch.Size([50]), type: torch.int32
target_lengths: torch.Size([2]), type: torch.int32
res.grad_fn: <MeanBackward1 object at 0x0000025E35A2CF98>
res.grad_fn: MeanBackward1

Process finished with exit code 0

loss.grad_fn's name not includes “cudnn”
Is there any problem with the code?