The problem is all_rays_cos is not on the gpu. When you define it here:
    all_rays_cos = torch.zeros(SB, NS, H*W)

you do not convert it to cuda. Just change that to this
    all_rays_cos = torch.zeros(SB, NS, H*W).to(device)