diff --git a/README.md b/README.md
index e63e346..b6ccb21 100644
--- a/README.md
+++ b/README.md
@@ -1,11 +1,13 @@
-# Optimal Transport Kernel
+# Optimal Transport Kernel Embedding
 
-The repository implements the Optimal Transport Kernel (OTK) described in the following paper
+This repository implements the Optimal Transport Kernel Embedding (OTKE) described in the following paper
 
 >Grégoire Mialon*, Dexiong Chen*, Alexandre d'Aspremont, Julien Mairal.
-[An Optimal Transport Kernel for Feature Aggregation and its Relationship to Attention][1]. preprint arXiv, 2020.
+[A Trainable Optimal Transport Embedding for Feature Aggregation and its Relationship to Attention][1]. ICLR 2021.
 <br/>*Equal contribution
 
+TLDR; Our paper demonstrates the advantage of our embedding over usual aggregation method (e.g, mean pooling, max pooling or attention) when faced to data composed of large sets of features, such as biological sequences, sentences or even images. Our embedding can be learned either with or without labels, and used alone as a kernel method or as a layer in larger models.
+
 ## A short description about the module
 
 The principal module is implemented in `otk/layers.py` as `OTKernel`. It is generally used with a non-linear layer. Combined with the non-linear layer, it takes a sequence or image tensor as input, and performs a non-linear embedding and an adaptive pooling (attention + pooling) based on optimal transport. Specifically, given a sequence x as input, it first computes the optimal transport plan from x to some reference z (left figure). The optimal transport plan, interpreted as the attention score, is then used to obtain a new sequence of the same size as z following a non-linear mapping (right figure). See more details in our [paper][1].
diff --git a/experiments/baseline_approxrepset.py b/experiments/baseline_approxrepset.py
new file mode 100644
index 0000000..7253444
--- /dev/null
+++ b/experiments/baseline_approxrepset.py
@@ -0,0 +1,254 @@
+import argparse
+import copy
+from math import ceil
+import numpy as np
+import os
+import torch
+import torch.nn.functional as F
+
+from loaders import load_data, load_masks
+from repset.approxrepset.models import ApproxRepSet
+from repset.approxrepset.utils import accuracy, AverageMeter
+
+"""
+This is the ApproxRepSet baseline using the code from the paper
+http://proceedings.mlr.press/v108/skianis20a/skianis20a.pdf.
+"""
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+def load_args():
+    parser = argparse.ArgumentParser(
+        description="baseline approxrepset for SST-2",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        '--dataset', type=str, default='sst-2_bert_mask', choices=['sst-2_bert_mask', 'sst-2_proto']
+    )
+    parser.add_argument(
+        '--seed', type=int, default=1, help='random seed')
+    parser.add_argument(
+        '--batch-size', type=int, default=64,
+        help='input batch size for training')
+    parser.add_argument(
+        '--epochs', type=int, default=30, metavar='N',
+        help='number of epochs to train')
+    parser.add_argument(
+        '--lr', type=float, default=1e-3, help='learning rate for the optimizer')
+    parser.add_argument(
+        '--heads', type=int, default=10, help='number of heads for attention layer')
+    parser.add_argument(
+        '--out-size', type=int, default=20, help='number of supports for attention layer')
+    parser.add_argument(
+        '--dim-hidden', type=int, default=768, help='dimension of each vector')
+    parser.add_argument(
+        "--outdir", default="results/", type=str, help="output path")
+    args = parser.parse_args()
+    args.use_cuda = torch.cuda.is_available()
+    # check shape
+
+    args.save_logs = False
+    if args.outdir != "":
+        args.save_logs = True
+        outdir = args.outdir + args.dataset
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + "/sup"
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + f"/approxrepset"
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + '/learning_{}_{}_{}'.format(
+            args.batch_size, args.epochs, args.lr)
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + '/refset_{}_{}_{}'.format(
+            args.heads, args.out_size, args.dim_hidden)
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        args.outdir = outdir
+
+    return args
+
+def main():
+    args = load_args()
+    print(args)
+    torch.manual_seed(args.seed)
+    if args.use_cuda:
+        torch.cuda.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    errs = list()
+
+    X_train, y_train, X_test, y_test, _ = load_data(dataset=args.dataset)
+    mask_train, mask_test, _ = load_masks(args.dataset)
+
+    X_train *= mask_train
+    X_test *= mask_test
+
+    X_train = X_train.permute(0, 2, 1).numpy()
+    X_test = X_test.permute(0, 2, 1).numpy()
+
+    y_train_ = np.zeros((y_train.size, y_train.max()+1)) # added that
+    y_train_[np.arange(y_train.size),y_train] = 1 # added that
+    y_train = y_train_
+
+    y_test_ = np.zeros((y_test.size, y_test.max()+1)) # added that
+    y_test_[np.arange(y_test.size),y_test] = 1 # added that
+    y_test = y_test_
+
+    n_train = int(0.8 * X_train.shape[0])
+    n_val = X_train.shape[0] - n_train
+    # n_train = y_train.shape[0]
+    n_test = y_test.shape[0]
+
+    idx = np.random.permutation(n_train)
+    n_train_batches = ceil(n_train / args.batch_size)
+    train_batches = list()
+
+    for i in range(n_train_batches):
+        max_card = max([X_train[idx[j]].shape[1] for j in range(
+            i * args.batch_size,min(( i+ 1) * args.batch_size, n_train))])
+        X = np.zeros((min((i + 1) * args.batch_size, n_train) - i * args.batch_size,
+                      max_card, args.dim_hidden))
+        for j in range(i * args.batch_size, min((i + 1) * args.batch_size, n_train)):
+            X[j - i * args.batch_size, :X_train[idx[j]].shape[1], :] = X_train[idx[j]].T
+        X = torch.FloatTensor(X).to(device)
+        y = torch.LongTensor(np.where(y_train[idx[i * args.batch_size:min((i + 1) * args.batch_size, n_train)]])[1]).to(device)
+        train_batches.append((X, y))
+
+    idx = np.random.permutation(range(n_train, X_train.shape[0]))
+    n_val_batches = ceil(n_val / args.batch_size)
+    val_batches = list()
+
+    for i in range(n_val_batches):
+        max_card = max([X_train[idx[j]].shape[1] for j in range(
+            i * args.batch_size,min(( i+ 1) * args.batch_size, n_val))])
+        X = np.zeros((min((i + 1) * args.batch_size, n_val) - i * args.batch_size,
+                      max_card, args.dim_hidden))
+        for j in range(i * args.batch_size, min((i + 1) * args.batch_size, n_val)):
+            X[j - i * args.batch_size, :X_train[idx[j]].shape[1], :] = X_train[idx[j]].T
+        X = torch.FloatTensor(X).to(device)
+        y = torch.LongTensor(np.where(y_train[idx[i * args.batch_size:min((i + 1) * args.batch_size, n_val)]])[1]).to(device)
+        val_batches.append((X, y))
+
+    n_test_batches = ceil(n_test / args.batch_size)
+    test_batches = list()
+
+    for i in range(n_test_batches):
+        max_card = max([X_test[j].shape[1] for j in range(
+            i * args.batch_size, min((i+1) * args.batch_size, n_test))])
+        X = np.zeros((min((i + 1) * args.batch_size, n_test) - i * args.batch_size,
+                      max_card, args.dim_hidden))
+        for j in range(i * args.batch_size, min((i + 1) * args.batch_size, n_test)):
+            X[j - i * args.batch_size, :X_test[j].shape[1], :] = X_test[j].T
+        X = torch.FloatTensor(X).to(device)
+        y = torch.LongTensor(np.where(
+            y_test[i * args.batch_size:min((i + 1) * args.batch_size, n_test)])[1]).to(device)
+        test_batches.append((X, y))
+
+    model = ApproxRepSet(args.heads, args.out_size, args.dim_hidden,
+                         n_classes=2, device=device).to(device)
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    def train(X, y):
+        optimizer.zero_grad()
+        output = model(X)
+        loss_train = F.cross_entropy(output, y)
+        loss_train.backward()
+        optimizer.step()
+        return output, loss_train
+
+    def test(X, y):
+        output = model(X)
+        loss_test = F.cross_entropy(output, y)
+        return output, loss_test
+
+    best_loss = float('inf')
+    for epoch in range(args.epochs):
+
+        model.train()
+        train_loss = AverageMeter()
+        train_err = AverageMeter()
+
+        for X, y in train_batches:
+            output, loss = train(X, y)
+
+            train_loss.update(loss.item(), output.size(0))
+            train_err.update(1 - accuracy(output.data, y.data), output.size(0))
+        
+        model.eval()
+        val_loss = AverageMeter()
+        val_acc = AverageMeter()
+
+        for X, y in val_batches:
+            output, loss = test(X, y)
+            val_loss.update(loss.item(), output.size(0))
+            val_acc.update(accuracy(output.data, y.data), output.size(0))
+        
+        if val_loss.avg < best_loss:
+            best_loss = val_loss.avg
+            best_acc = val_acc.avg
+            best_epoch = epoch + 1
+            best_weights = copy.deepcopy(model.state_dict())
+
+        print("epoch:", '%03d' % (epoch+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+            "train_acc=", "{:.5f}".format(1 - train_err.avg), "val_loss= {}".format(val_loss.avg),
+            "val_acc= {}".format(val_acc.avg))
+
+    model.load_state_dict(best_weights)
+    print("Testing...")
+    model.eval()
+
+    test_loss = AverageMeter()
+    test_err = AverageMeter()
+
+    for X, y in test_batches:
+        output, loss = test(X, y)
+        
+        test_loss.update(loss.item(), output.size(0))
+        test_err.update(1 - accuracy(output.data, y.data), output.size(0))
+
+    print("train_loss=", "{:.5f}".format(train_loss.avg),
+          "train_acc=", "{:.5f}".format(1 - train_err.avg),
+          "test_loss=", "{:.5f}".format(test_loss.avg),
+          "test_acc=", "{:.5f}".format(1 - test_err.avg),
+          "best_epoch", "{:.5f}".format(best_epoch)
+          )
+    print()
+
+    errs.append(test_err.avg.cpu())
+
+    print("Average accuracy:", "{:.5f}".format(1 - np.mean(errs)))
+
+    if args.save_logs:
+        print('Saving logs...')
+        data = {
+            'score': 1 - test_err.avg.cpu(),
+            'best_epoch': best_epoch,
+            'best_loss': best_loss,
+            'train_acc': 1 - train_err.avg.cpu(),
+            'val_score': best_acc.cpu(),
+            'args': args
+            }
+        np.save(os.path.join(args.outdir, f"seed_{args.seed}_results.npy"),
+                data)
+    return
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/experiments/baseline_repset.py b/experiments/baseline_repset.py
new file mode 100644
index 0000000..ca52a2a
--- /dev/null
+++ b/experiments/baseline_repset.py
@@ -0,0 +1,207 @@
+import argparse
+import numpy as np
+import os
+import torch
+import copy
+
+from math import ceil
+from sklearn.metrics import accuracy_score,log_loss
+from timeit import default_timer as timer
+
+from loaders import load_data, load_masks
+from repset.repset.models import RepSet
+from repset.repset.utils import AverageMeter
+
+"""
+This is the RepSet baseline using the code from the paper
+http://proceedings.mlr.press/v108/skianis20a/skianis20a.pdf.
+"""
+
+def load_args():
+    parser = argparse.ArgumentParser(
+        description="baseline repset for SST-2",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    parser.add_argument(
+        '--dataset', type=str, default='sst-2_bert_mask', choices=['sst-2_bert_mask', 'sst-2_proto']
+    )
+    parser.add_argument(
+        '--seed', type=int, default=1, help='random seed')
+    parser.add_argument(
+        '--batch-size', type=int, default=64,
+        help='input batch size for training')
+    parser.add_argument(
+        '--epochs', type=int, default=10, metavar='N',
+        help='number of epochs to train')
+    parser.add_argument(
+        '--heads', type=int, default=10, help='number of heads for attention layer')
+    parser.add_argument(
+        '--out-size', type=int, default=20, help='number of supports for attention layer')
+    parser.add_argument(
+        '--dim-hidden', type=int, default=768, help='dimension of each vector')
+    parser.add_argument(
+        "--outdir", default="results/", type=str, help="output path")
+    parser.add_argument(
+        "--lr", type=float, default=0.01, help='initial learning rate')
+    args = parser.parse_args()
+    args.use_cuda = torch.cuda.is_available()
+    # check shape
+
+    args.save_logs = False
+    if args.outdir != "":
+        args.save_logs = True
+        outdir = args.outdir + args.dataset
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + "/sup"
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + f"/repset"
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir+'/learning_{}_{}_{}'.format(
+            args.batch_size, args.epochs, args.lr)
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        outdir = outdir + '/refset_{}_{}_{}'.format(
+            args.heads, args.out_size, args.dim_hidden)
+        if not os.path.exists(outdir):
+            try:
+                os.makedirs(outdir)
+            except:
+                pass
+        args.outdir = outdir
+
+    return args
+
+def main():
+    args = load_args()
+    print(args)
+    np.random.seed(args.seed)
+    errs = list()
+
+    X_train, y_train, X_test, y_test, _ = load_data(dataset=args.dataset)
+    mask_train, mask_test, _ = load_masks(args.dataset)
+
+    X_train *= mask_train
+    X_test *= mask_test
+
+    X_train = X_train.permute(0, 2, 1).numpy()
+    X_test = X_test.permute(0, 2, 1).numpy()
+
+    n_train = int(0.8 * X_train.shape[0])
+    n_val = X_train.shape[0] - n_train
+    # n_train = y_train.shape[0]
+    n_test = y_test.shape[0]
+
+    idx = np.random.permutation(n_train)
+    n_train_batches = ceil(n_train/args.batch_size)
+    train_batches = list()
+
+    for i in range(n_train_batches):
+        train_batches.append((X_train[idx[i * args.batch_size:min((i+1) * args.batch_size, n_train)]],
+                              y_train[idx[i * args.batch_size:min((i+1) * args.batch_size, n_train)]]))
+
+    idx = np.random.permutation(range(n_train, X_train.shape[0]))
+    n_val_batches = ceil(n_val / args.batch_size)
+    val_batches = list()
+
+    for i in range(n_val_batches):
+        val_batches.append((X_train[idx[i * args.batch_size:min((i+1) * args.batch_size, n_val)]],
+                              y_train[idx[i * args.batch_size:min((i+1) * args.batch_size, n_val)]]))
+
+    n_test_batches = ceil(n_test/args.batch_size)
+    test_batches = list()
+    for i in range(n_test_batches):
+        test_batches.append((X_test[i*args.batch_size:min((i+1)*args.batch_size, n_test)],
+                             y_test[i*args.batch_size:min((i+1)*args.batch_size, n_test)]))
+
+    model = RepSet(args.lr, args.heads, args.out_size, args.dim_hidden, n_classes=2)
+
+    best_loss = float('inf')
+    for epoch in range(args.epochs):
+
+        train_loss = AverageMeter()
+        train_err = AverageMeter()
+        
+        for X, y in train_batches:
+            y_ = np.zeros((y.size, y.max()+1)) # added that
+            y_[np.arange(y.size),y] = 1 # added that
+            y_pred = model.train(X, y_)
+            train_loss.update(log_loss(y_, y_pred), n_train)
+            train_err.update(1-accuracy_score(np.argmax(y_, axis=1), np.argmax(y_pred, axis=1)), y_.shape[0])
+
+        val_loss = AverageMeter()
+        val_acc = AverageMeter()
+
+        for X, y in val_batches:
+            y_ = np.zeros((y.size, y.max()+1)) # added that
+            y_[np.arange(y.size),y] = 1 # added that
+            y_pred = model.test(X)
+            val_loss.update(log_loss(y_, y_pred), n_val)
+            val_acc.update(accuracy_score(np.argmax(y_, axis=1), np.argmax(y_pred, axis=1)), y_.shape[0])
+
+        if val_loss.avg < best_loss:
+            best_loss = val_loss.avg
+            best_acc = val_acc.avg
+            best_epoch = epoch + 1
+            best_model = copy.deepcopy(model)
+
+        print("epoch:", '%03d' % (epoch+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+            "train_acc=", "{:.5f}".format(1 - train_err.avg), "val_loss= {}".format(val_loss.avg),
+            "val_acc= {}".format(val_acc.avg)
+            )
+
+    print("Testing...")
+    test_loss = AverageMeter()
+    test_err = AverageMeter()
+
+    for X, y in test_batches:
+        y_ = np.zeros((y.size, y.max()+1)) # added that
+        y_[np.arange(y.size),y] = 1 # added that
+        y_test_ = np.zeros((y_test.size, y_test.max()+1)) # added that
+        y_test_[np.arange(y_test.size),y_test] = 1 # added that
+        y_pred = best_model.test(X)
+
+        test_loss.update(log_loss(y_, y_pred), y_test_.size)
+        test_err.update(1-accuracy_score(np.argmax(y_, axis=1), np.argmax(y_pred, axis=1)), y_.shape[0])
+
+    print("train_loss=", "{:.5f}".format(train_loss.avg),
+          "train_acc=", "{:.5f}".format(1 - train_err.avg),
+          "test_loss=", "{:.5f}".format(test_loss.avg),
+          "test_acc=", "{:.5f}".format(1 - test_err.avg),
+          "best_epoch=", "{:.5f}".format(best_epoch) 
+        )
+    print()
+
+    errs.append(test_err.avg)
+
+    print("Average accuracy:", "{:.5f}".format(1 - np.mean(errs)))
+
+    if args.save_logs:
+        print('Saving logs...')
+        data = {
+            'score': 1 - test_err.avg,
+            'best_epoch': best_epoch,
+            'best_loss': best_loss,
+            'train_acc': 1 - train_err.avg,
+            'val_score': best_acc,
+            'args': args
+            }
+        np.save(os.path.join(args.outdir, f"seed_{args.seed}_results.npy"),
+                data)
+    return
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/otk/sinkhorn.py b/otk/sinkhorn.py
index 581f7ae..2723411 100644
--- a/otk/sinkhorn.py
+++ b/otk/sinkhorn.py
@@ -154,7 +154,7 @@ def wasserstein_kmeans(x, n_clusters, out_size, eps=1.0, block_size=None, max_it
         del x_batch
         sim = wass_sim.mean()
         if verbose and (n_iter + 1) % 10 == 0:
-            print("Wassestein spherical kmeans iter {}, objective value {}".format(
+            print("Wasserstein spherical kmeans iter {}, objective value {}".format(
                   n_iter + 1, sim))
 
         for j in range(n_clusters):
diff --git a/repset/.gitignore b/repset/.gitignore
new file mode 100644
index 0000000..894a44c
--- /dev/null
+++ b/repset/.gitignore
@@ -0,0 +1,104 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
diff --git a/repset/README.md b/repset/README.md
new file mode 100644
index 0000000..617167d
--- /dev/null
+++ b/repset/README.md
@@ -0,0 +1,39 @@
+## Rep the Set: Neural Networks for Learning Set Representations
+Code for the paper [Rep the Set: Neural Networks for Learning Set Representations](http://proceedings.mlr.press/v108/skianis20a/skianis20a.pdf).
+
+### Requirements
+Code is written in Python 3.6 and requires:
+* lap 0.4.0
+* PyTorch 1.1
+* scikit-learn 0.21
+
+### Run the model
+You can specify the dataset and the hyperparameters in the main.py files.
+
+For the exact model, run:
+
+```
+python repset/main.py
+```
+
+For the approximate model, run:
+
+```
+python approxrepset/main.py
+```
+
+### Cite
+Please cite our paper if you use this code:
+```
+@inproceedings{skianis2020rep,
+  title={Rep the Set: Neural Networks for Learning Set Representations},
+  author={Skianis, Konstantinos and Nikolentzos, Giannis and Limnios, Stratis and Vazirgiannis, Michalis},
+  booktitle={Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics},
+  pages={1410--1420},
+  year={2020}
+}
+```
+
+-----------
+
+Provided for academic use only
diff --git a/repset/approxrepset/main.py b/repset/approxrepset/main.py
new file mode 100644
index 0000000..5368751
--- /dev/null
+++ b/repset/approxrepset/main.py
@@ -0,0 +1,101 @@
+import torch
+import numpy as np
+import torch.nn.functional as F
+from math import ceil
+from models import ApproxRepSet
+from utils import load_data,accuracy,AverageMeter
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Parameters
+dataset = 'twitter'
+epochs = 30 # number of iterations
+lr = 1e-3 # learning rate
+n_hidden_sets = 10 # number of hidden sets
+n_elements = 20 # cardinality of each hidden set
+d = 300 # dimension of each vector
+batch_size = 64  # batch size
+cv_folds = 5 # number of folds for cross-validation
+
+errs = list()
+
+for it in range(cv_folds):
+    X_train, X_test, y_train, y_test, n_classes, _, _, _ , _ = load_data(dataset, it)
+
+    n_train = y_train.shape[0]
+    n_test = y_test.shape[0]
+
+    idx = np.random.permutation(n_train)
+    n_train_batches = ceil(n_train/batch_size)
+    train_batches = list()
+    for i in range(n_train_batches):
+        max_card = max([X_train[idx[j]].shape[1] for j in range(i*batch_size,min((i+1)*batch_size, n_train))])
+        X = np.zeros((min((i+1)*batch_size, n_train)-i*batch_size, max_card, d))
+        for j in range(i*batch_size,min((i+1)*batch_size, n_train)):
+            X[j-i*batch_size,:X_train[idx[j]].shape[1],:] = X_train[idx[j]].T
+        X = torch.FloatTensor(X).to(device)
+        y = torch.LongTensor(np.where(y_train[idx[i*batch_size:min((i+1)*batch_size, n_train)]])[1]).to(device)
+        train_batches.append((X, y))
+
+    n_test_batches = ceil(n_test/batch_size)
+    test_batches = list()
+    for i in range(n_test_batches):
+        max_card = max([X_test[j].shape[1] for j in range(i*batch_size,min((i+1)*batch_size, n_test))])
+        X = np.zeros((min((i+1)*batch_size, n_test)-i*batch_size, max_card, d))
+        for j in range(i*batch_size,min((i+1)*batch_size, n_test)):
+            X[j-i*batch_size,:X_test[j].shape[1],:] = X_test[j].T
+        X = torch.FloatTensor(X).to(device)
+        y = torch.LongTensor(np.where(y_test[i*batch_size:min((i+1)*batch_size, n_test)])[1]).to(device)
+        test_batches.append((X, y))
+
+    model = ApproxRepSet(n_hidden_sets, n_elements, d, n_classes, device).to(device)
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+ 
+    def train(X, y):
+        optimizer.zero_grad()
+        output = model(X)
+        loss_train = F.cross_entropy(output, y)
+        loss_train.backward()
+        optimizer.step()
+        return output, loss_train
+
+    def test(X, y):
+        output = model(X)
+        loss_test = F.cross_entropy(output, y)
+        return output, loss_test
+
+    model.train()
+    for epoch in range(epochs):
+
+        train_loss = AverageMeter()
+        train_err = AverageMeter()
+
+        for X, y in train_batches:
+            output, loss = train(X, y)
+
+            train_loss.update(loss.item(), output.size(0))
+            train_err.update(1-accuracy(output.data, y.data), output.size(0))
+
+        print("Cross-val iter:", '%02d' % (it+1), "epoch:", '%03d' % (epoch+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+            "train_err=", "{:.5f}".format(train_err.avg))
+
+    model.eval()
+
+    test_loss = AverageMeter()
+    test_err = AverageMeter()
+
+    for X, y in test_batches:
+        output, loss = test(X, y)
+        
+        test_loss.update(loss.item(), output.size(0))
+        test_err.update(1-accuracy(output.data, y.data), output.size(0))
+
+    print("Cross-val iter:", '%02d' % (it+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+        "train_err=", "{:.5f}".format(train_err.avg), "test_loss=", "{:.5f}".format(test_loss.avg), "test_err=", "{:.5f}".format(test_err.avg))
+    print()
+
+    errs.append(test_err.avg.cpu())
+
+print("Average error:", "{:.5f}".format(np.mean(errs)))
+print("Standard deviation:", "{:.5f}".format(np.std(errs)))
\ No newline at end of file
diff --git a/repset/approxrepset/models.py b/repset/approxrepset/models.py
new file mode 100644
index 0000000..3f41a72
--- /dev/null
+++ b/repset/approxrepset/models.py
@@ -0,0 +1,31 @@
+import torch
+import torch.nn as nn
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F
+
+class ApproxRepSet(torch.nn.Module):
+
+    def __init__(self, n_hidden_sets, n_elements, d, n_classes, device):
+        super(ApproxRepSet, self).__init__()
+        self.n_hidden_sets = n_hidden_sets
+        self.n_elements = n_elements
+        
+        self.Wc = Parameter(torch.FloatTensor(d, n_hidden_sets*n_elements))
+        self.fc1 = nn.Linear(n_hidden_sets, 32)
+        self.fc2 = nn.Linear(32, n_classes)
+        self.relu = nn.ReLU()
+
+        self.init_weights()
+
+    def init_weights(self):
+        self.Wc.data.uniform_(-1, 1)
+
+    def forward(self, X):
+        t = self.relu(torch.matmul(X, self.Wc))
+        t = t.view(t.size()[0], t.size()[1], self.n_elements, self.n_hidden_sets)
+        t,_ = torch.max(t, dim=2)
+        t = torch.sum(t, dim=1)
+        t = self.relu(self.fc1(t))
+        out = self.fc2(t)
+
+        return F.log_softmax(out, dim=1)
diff --git a/repset/approxrepset/utils.py b/repset/approxrepset/utils.py
new file mode 100644
index 0000000..a625bca
--- /dev/null
+++ b/repset/approxrepset/utils.py
@@ -0,0 +1,56 @@
+import numpy as np
+from sklearn.preprocessing import LabelEncoder
+from scipy.io import loadmat
+
+def load_data(dataset, cv_fold):
+    data = loadmat('../datasets/'+dataset+'_tr_te_split.mat')
+    X = data['X']
+    TR = data['TR'][cv_fold,:]-1
+    TE = data['TE'][cv_fold,:]-1
+
+    X_train = X[:,TR][0]
+    X_test = X[:,TE][0]
+
+    BOW_X = data['BOW_X']
+    BOW_X_train = BOW_X[:,TR]
+    BOW_X_test = BOW_X[:,TE]
+
+    class_labels = np.array(data['Y'][0])
+    le = LabelEncoder()
+    class_labels = le.fit_transform(class_labels)
+    n_classes = np.unique(class_labels).size
+    y = np.zeros((class_labels.size, n_classes))
+    for i in range(class_labels.size):
+        y[i,class_labels[i]] = 1
+    y_train = y[TR,:]
+    y_test = y[TE,:]
+
+    words = data['words']
+    words_train = words[:,TR]
+    words_test = words[:,TE]
+
+    return X_train, X_test, y_train, y_test, n_classes, BOW_X_train, BOW_X_test, words_train, words_test
+
+def accuracy(output, labels):
+    preds = output.max(1)[1].type_as(labels)
+    correct = preds.eq(labels).double()
+    correct = correct.sum()
+    return correct / len(labels)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
\ No newline at end of file
diff --git a/repset/repset/main.py b/repset/repset/main.py
new file mode 100644
index 0000000..c52754c
--- /dev/null
+++ b/repset/repset/main.py
@@ -0,0 +1,69 @@
+import numpy as np
+from math import ceil
+from sklearn.metrics import accuracy_score,log_loss
+from models import RepSet
+from utils import load_data,AverageMeter
+
+# Parameters
+dataset = 'twitter'
+epochs = 10 # number of iterations
+lr = 0.01 # learning rate
+n_hidden_sets = 10 # number of hidden sets
+n_elements = 20 # cardinality of each hidden set
+d = 300 # dimension of each vector
+batch_size = 64  # batch size
+cv_folds = 5 # number of folds for cross-validation
+
+errs = list()
+
+for it in range(cv_folds):
+    X_train, X_test, y_train, y_test, n_classes, _, _, _ , _ = load_data(dataset, it)
+
+
+    n_train = y_train.shape[0]
+    n_test = y_test.shape[0]
+
+    idx = np.random.permutation(n_train)
+    n_train_batches = ceil(n_train/batch_size)
+    train_batches = list()
+
+    for i in range(n_train_batches):
+        train_batches.append((X_train[idx[i*batch_size:min((i+1)*batch_size, n_train)]], y_train[idx[i*batch_size:min((i+1)*batch_size, n_train)]]))
+
+    n_test_batches = ceil(n_test/batch_size)
+    test_batches = list()
+    for i in range(n_test_batches):
+        test_batches.append((X_test[i*batch_size:min((i+1)*batch_size, n_test)], y_test[i*batch_size:min((i+1)*batch_size, n_test)]))
+
+    model = RepSet(lr, n_hidden_sets, n_elements, d, n_classes)
+ 
+    for epoch in range(epochs):
+
+        train_loss = AverageMeter()
+        train_err = AverageMeter()
+
+        for X, y in train_batches:
+            y_pred = model.train(X, y)
+
+            train_loss.update(log_loss(y, y_pred), y_train.size)
+            train_err.update(1-accuracy_score(np.argmax(y, axis=1), np.argmax(y_pred, axis=1)), y.shape[0])
+
+        print("Cross-val iter:", '%02d' % (it+1), "epoch:", '%03d' % (epoch+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+            "train_err=", "{:.5f}".format(train_err.avg))
+
+    test_loss = AverageMeter()
+    test_err = AverageMeter()
+
+    for X, y in test_batches:
+        y_pred = model.test(X)
+        test_loss.update(log_loss(y, y_pred), y_test.size)
+        test_err.update(1-accuracy_score(np.argmax(y, axis=1), np.argmax(y_pred, axis=1)), y.shape[0])
+
+    print("Cross-val iter:", '%02d' % (it+1), "train_loss=", "{:.5f}".format(train_loss.avg),
+        "train_err=", "{:.5f}".format(train_err.avg), "test_loss=", "{:.5f}".format(test_loss.avg), "test_err=", "{:.5f}".format(test_err.avg))
+    print()
+
+    errs.append(test_err.avg)
+
+print("Average error:", "{:.5f}".format(np.mean(errs)))
+print("Standard deviation:", "{:.5f}".format(np.std(errs)))
diff --git a/repset/repset/models.py b/repset/repset/models.py
new file mode 100644
index 0000000..63bf2d1
--- /dev/null
+++ b/repset/repset/models.py
@@ -0,0 +1,111 @@
+import numpy as np
+from lap import lapjv
+
+class RepSet:
+    def __init__(self, lr, n_hidden_sets, n_elements, d, n_classes):
+        self.lr = lr
+        self.n_hidden_sets = n_hidden_sets
+        self.n_elements = n_elements
+        self.d = d
+        self.n_classes = n_classes
+
+        self.t = 0
+        self.beta_1 = 0.9
+        self.beta_2 = 0.999
+        self.epsilon = 1e-8
+        self.m_t = [0]*n_hidden_sets
+        self.v_t = [0]*n_hidden_sets
+        self.m_t_c = 0
+        self.v_t_c = 0
+
+        self.Ws = np.random.randn(n_hidden_sets, n_elements, d)
+        self.Wc = np.random.randn(n_hidden_sets+1, n_classes)
+
+
+    def train(self, X, y):
+        # R = np.zeros((X.size, self.n_hidden_sets+1))
+        R = np.zeros((X.shape[0], self.n_hidden_sets+1))
+        R[:,-1] = 1
+
+        Ds = list()
+
+        for i in range(X.shape[0]):
+            Ds.append(list())
+
+            x = X[i]
+    
+            for j in range(self.n_hidden_sets):
+                W = self.Ws[j]
+                K = np.dot(W, x)
+                K[K<0] = 0
+
+                cost, x_lap, _ = lapjv(-K, extend_cost=True)
+
+                D = np.zeros((self.n_elements, x.shape[1]))
+                for k in range(self.n_elements):
+                    if x_lap[k] != -1:
+                        D[k, x_lap[k]] = 1
+
+                Ds[i].append(D)
+
+                cost_norm = cost/x.shape[1]
+                R[i,j] = -cost_norm
+
+        S = np.dot(R, self.Wc)
+        y_pred = np.exp(S)/np.sum(np.exp(S), axis=1).reshape(-1, 1)
+
+        E = y - y_pred
+
+        ## Backprop
+        upd_Ws = np.zeros((self.n_hidden_sets, self.n_elements, self.d))
+        upd_Wc = np.zeros((self.n_hidden_sets+1, self.n_classes))
+
+        for i in range(X.shape[0]):
+            x = X[i]
+    
+            for j in range(self.n_hidden_sets):
+                upd_Ws[j] = upd_Ws[j] + np.dot(Ds[i][j], x.T)*np.dot(E[i,:], self.Wc[j,:])
+
+            upd_Wc += np.outer(R[i,:].T, E[i,:])
+        
+        
+        self.t += 1
+        for j in range(self.n_hidden_sets):
+            g_t = upd_Ws[j]*1./x.shape[1]
+            self.m_t[j] = self.beta_1*self.m_t[j] + (1-self.beta_1)*g_t
+            self.v_t[j] = self.beta_2*self.v_t[j] + (1-self.beta_2)*(np.square(g_t))
+            m_cap = self.m_t[j]/(1-(self.beta_1**self.t))
+            v_cap = self.v_t[j]/(1-(self.beta_2**self.t))
+            self.Ws[j] = self.Ws[j] + (self.lr*m_cap)/(np.sqrt(v_cap)+self.epsilon)
+            
+        g_t = upd_Wc*1./x.shape[1]
+        self.m_t_c = self.beta_1*self.m_t_c + (1-self.beta_1)*g_t
+        self.v_t_c = self.beta_2*self.v_t_c + (1-self.beta_2)*np.square(g_t)
+        m_cap= self.m_t_c/(1-(self.beta_1**self.t))
+        v_cap = self.v_t_c/(1-(self.beta_2**self.t))
+        self.Wc = self.Wc + (self.lr*m_cap)/(np.sqrt(v_cap)+self.epsilon)
+
+        return y_pred
+
+
+    def test(self, X):
+        # R = np.zeros((X.size, self.n_hidden_sets+1))
+        R = np.zeros((X.shape[0], self.n_hidden_sets+1))
+        R[:,-1] = 1
+
+        for i in range(X.shape[0]):
+            x = X[i]
+    
+            for j in range(self.n_hidden_sets):
+                W = self.Ws[j]
+                K = np.dot(W, x)
+                K[K<0] = 0
+
+                cost, x_lap, _ = lapjv(-K, extend_cost=True)
+                cost_norm = cost/x.shape[1]
+                R[i,j] = -cost_norm
+
+        R = np.dot(R, self.Wc)
+        y_pred = np.exp(R)/np.sum(np.exp(R), axis=1).reshape(-1, 1)
+        
+        return y_pred
\ No newline at end of file
diff --git a/repset/repset/utils.py b/repset/repset/utils.py
new file mode 100644
index 0000000..07cf1be
--- /dev/null
+++ b/repset/repset/utils.py
@@ -0,0 +1,50 @@
+import numpy as np
+from sklearn.preprocessing import LabelEncoder
+from scipy.io import loadmat
+
+def load_data(dataset, cv_fold):
+    data = loadmat('../datasets/'+dataset+'_tr_te_split.mat')
+    X = data['X']
+    TR = data['TR'][cv_fold,:]-1
+    TE = data['TE'][cv_fold,:]-1
+
+    X_train = X[:,TR][0]
+    X_test = X[:,TE][0]
+
+    BOW_X = data['BOW_X']
+    BOW_X_train = BOW_X[:,TR]
+    BOW_X_test = BOW_X[:,TE]
+
+    class_labels = np.array(data['Y'][0])
+    le = LabelEncoder()
+    class_labels = le.fit_transform(class_labels)
+    n_classes = np.unique(class_labels).size
+    y = np.zeros((class_labels.size, n_classes))
+    for i in range(class_labels.size):
+        y[i,class_labels[i]] = 1
+    y_train = y[TR,:]
+    y_test = y[TE,:]
+
+    words = data['words']
+    words_train = words[:,TR]
+    words_test = words[:,TE]
+
+    return X_train, X_test, y_train, y_test, n_classes, BOW_X_train, BOW_X_test, words_train, words_test
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
\ No newline at end of file