Jinwee0125/iris_knn

Usage

import torch
import datasets

dataset = load_dataset("scikit-learn/iris")["train"].class_encode_column("Species").train_test_split(
        test_size=0.3,
        seed=42,
        stratify_by_column="Species"
)
features = ["SepalLengthCm","SepalWidthCm","PetalLengthCm","PetalWidthCm"]
batch = dataset["test"].with_format("torch")[:]
x = torch.stack([batch[col] for col in features],dim=1)
model = torch.jit.load("iris_knn.pt")
preds = model(x)