4cbe889c27f5a0c272c54b4e1065ef3525c2fc09,s2cnn/ops/gpu/s2_mm.py,,s2_mm,#Any#Any#,47

Before Change

    cuda_kernel = _setup_s2mm_cuda_kernel(nbatch=nbatch, nspec=nspec, nfeature_in=nfeature_in, nfeature_out=nfeature_out)

    stream = cuda_utils.Stream(ptr=torch.cuda.current_stream().cuda_stream)
    output = torch.cuda.FloatTensor(nspec, nbatch, nfeature_out, 2)
    cuda_kernel(block=(cuda_utils.CUDA_NUM_THREADS, 1, 1),
                grid=(cuda_utils.get_blocks(nspec * nbatch * nfeature_out, 1024), 1, 1),
                args=[x.data_ptr(), y.data_ptr(), output.data_ptr()],
                stream=stream)

After Change

    :return:  [l * m * n, batch,      feature_out, complex]
    """
    assert x.is_cuda and x.dtype == torch.float32
    assert y.is_cuda and y.dtype == torch.float32
    assert y.size(3) == 2
    assert x.size(3) == 2
    nbatch = x.size(1)
    nfeature_in = x.size(2)

Instances