spark isotonic regression to torch

isotonic_reg.py

import torch
from torch import nn
from torchinterp1d import Interp1d

class InterpModule(nn.Module):
    def __init__(self, X, y):
        super().__init__()
        self.X = X
        self.y = y
        self.interp = Interp1d()


    def forward(self, input: torch.Tensor):
        return self.interp.forward(self.X, self.y, prediction)

to install

pip install git+git://github.com/aliutkus/torchinterp1d@master

note the latest one depends on torch 1.6