integrate FAISS with a PyTorch model

import torch
import faiss
import numpy as np

# Assume we have a pre-trained PyTorch model for image feature extraction
class FeatureExtractor(torch.nn.Module):
    # ... model definition ...

    def forward(self, x):
        return self.features(x)

model = FeatureExtractor()
model.eval()

# Extract features from your image dataset
image_features = []
for image in dataset:
    with torch.no_grad():
        features = model(image).cpu().numpy()
    image_features.append(features)

image_features = np.vstack(image_features)

# Create FAISS index
index = faiss.IndexFlatL2(image_features.shape[1])
index.add(image_features)

# Now you can use this index for similarity search