Right now I’m doing the following. Is there a faster way?
            discount_matrix = torch.triu(torch.ones(
                size=(N, N)))
            for i in range(1, N):
                tril_idx = torch.triu_indices(
                    row=N,
                    col=N,
                    offset=i)
                discount_matrix[tril_idx[0], tril_idx[1]] *= c