diff --git a/classification/finetune_test.py b/classification/finetune_test.py
index 9b195a3..01a630d 100644
--- a/classification/finetune_test.py
+++ b/classification/finetune_test.py
@@ -37,49 +37,31 @@ def args_parser():
 import wandb
 import timm
 
-
-def test_finetune(model, trainset, testset, epochs, lr):
-    model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
-    criterion = nn.CrossEntropyLoss()
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
-    model.train()
-    # epochs = 1
-    for ep in tqdm(range(epochs)):
-        for inputs, targets in tqdm(trainloader):
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-        # scheduler.step()
-    model.eval()
-    acc, test_loss = test(model, testloader, torch.device('cuda'))
-    return round(acc,2), round(test_loss,2)
-
 def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
     trainloader = DataLoader(trainset, batch_size=args.bs, shuffle=True, num_workers=4,drop_last=True)
     testloader = DataLoader(testset, batch_size=args.bs, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    # optimizer = optim.SGD(model.module.score.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    optimizer = optim.SGD(model.module.score.parameters(), lr=lr, weight_decay=0.01)
+
+
     criterion = nn.CrossEntropyLoss()
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
+    # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
     accs = []
     losses = []
-    # epochs = 1
-    for ep in tqdm(range(epochs)):
+    for ep in tqdm(range(epochs), desc="epoch"):
         model.train()
-        for inputs, targets in trainloader:
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
+        train_loss = 0
+        for batch in tqdm(trainloader, desc="batch"):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
+            outputs = model(inputs, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, targets)
+            train_loss += loss.item()
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
+        wandb.log({f'{mode}: train loss': train_loss / len(trainloader.dataset)})
         test_acc, test_loss = test(model, testloader, torch.device('cuda'))
         accs.append(test_acc)
         losses.append(test_loss)
@@ -89,13 +71,12 @@ def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
 
 if __name__ == '__main__':
     args = args_parser()
-    import wandb 
+    import wandb
     wandb.init(
-    project="sohpon classification finetune test",  
-    entity="sophon",
-    config = args,
-    name = f"{args.arch}_{args.dataset}" ,
-    notes = args.notes)   
+        project="sophon classification finetune test nlp",
+        config = args,
+        name = f"{args.arch}_{args.dataset}" ,
+        notes = args.notes)
     seed = args.seed
     set_seed(seed)
     trainset_tar, testset_tar = get_dataset(args.dataset, '../../../datasets', args=args)
@@ -113,9 +94,11 @@ def test_finetune_final(args, mode, model, trainset, testset, epochs, lr):
         acc, test_loss = test_finetune_final(args, 'normal pretrained/direct all', model.cuda(), trainset_tar, testset_tar, args.truly_finetune_epochs, args.finetune_lr)
 
     # ### train from scratch
-    elif args.start == 'sratch':
+    elif args.start == 'scratch':
         print('========test train from scratch=========')
         acc, test_loss = test_finetune_final(args, 'train from scratch/', model.cuda(), trainset_tar, testset_tar, args.truly_finetune_epochs, args.finetune_lr)
-    
+
     else:
-        assert(0)
\ No newline at end of file
+        assert(0)
+
+    print(f'test accuracy is {acc}, test loss is {test_loss}')
\ No newline at end of file
diff --git a/classification/inverse_loss.py b/classification/inverse_loss.py
index aee7e30..a85b198 100644
--- a/classification/inverse_loss.py
+++ b/classification/inverse_loss.py
@@ -6,6 +6,8 @@
 import argparse
 import json
 import sys
+from transformers import GPT2LMHeadModel
+
 sys.path.append('../')
 def args_parser():
     parser = argparse.ArgumentParser(description='train N shadow models')
@@ -19,7 +21,7 @@ def args_parser():
     parser.add_argument('--test_iterval', default=10, type=int)
     parser.add_argument('--arch', default='caformer', type=str)
     parser.add_argument('--gpus', default='0,1', type=str)
-    parser.add_argument('--dataset', default='', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC'])
+    parser.add_argument('--dataset', default='', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC', 'IMDB', 'PILE'])
     parser.add_argument('--finetune_epochs', default=1, type=int)
     parser.add_argument('--truly_finetune_epochs', default=20, type=int)
     parser.add_argument('--finetune_lr', default=0.0001, type=float)
@@ -52,9 +54,9 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
     test_loss = 0
     test_accuracy = 0
     total_test = 0
+
     for index,batch in enumerate(batches):
-        data, labels = batch
-        data, labels = data.to(device), labels.to(device)
+        data, labels, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
         adaptation_indices = np.zeros(data.size(0), dtype=bool)
         # adaptation_indices[np.arange(shots*ways)] = True
         adaptation_indices[np.random.choice(np.arange(data.size(0)), shots*ways, replace=False)] = True
@@ -65,13 +67,13 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
         current_test = evaluation_data.shape[0]
         # print(current_test)
         total_test += current_test
-        adaptation_error = loss(learner(adaptation_data), adaptation_labels)
+        adaptation_error = loss(learner(adaptation_data).logits, adaptation_labels)
         if index == 0:
-            current_grads = learner.adapt(adaptation_error,None) 
+            current_grads = learner.adapt(adaptation_error,None, allow_nograd=True) #allow_nograd?
         else:
             last_grads = current_grads
-            current_grads = learner.adapt(adaptation_error,last_grads) 
-        predictions = learner(evaluation_data)
+            current_grads = learner.adapt(adaptation_error,last_grads, allow_nograd=True) 
+        predictions = learner(evaluation_data).logits
         evaluation_error = loss(1-predictions, evaluation_labels)  
         evaluation_accuracy = accuracy(predictions, evaluation_labels)
         test_loss += evaluation_error*current_test
@@ -81,9 +83,9 @@ def fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
 
 def test_finetune(model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
@@ -101,19 +103,19 @@ def test_finetune(model, trainset, testset, epochs, lr):
 
 def test_finetune_final(mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4, drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4, drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
     # epochs = 1
     for ep in tqdm(range(epochs)):
         model.train()
-        for inputs, targets in tqdm(trainloader):
-            inputs, targets = inputs.cuda(), targets.cuda()  
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
+        for batch in tqdm(trainloader):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
+            outputs = model(inputs, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, targets)
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
@@ -140,10 +142,11 @@ def main(
     print("Hostname:", hostname)
     ip_address = socket.gethostbyname(hostname)
     args.from_machine = ip_address
+    lm_model = GPT2LMHeadModel.from_pretrained("gpt2").cuda()
 
     wandb.init(
     project="sophon classification",  
-    entity="sophon",
+    # entity="sophon",
     config = args,
     name = f"{args.dataset}_alpha{args.alpha}_beta{args.beta}_ml{args.ml_loop}_nl{args.nl_loop}_batches{args.adaptation_steps}" ,
     notes= args.notes,         
@@ -157,14 +160,14 @@ def main(
         # torch.cuda.manual_seed(seed)
         device = torch.device('cuda')
     wandb.log({'seed':seed})
-        save_path = args.root + '/inverse_loss'+ '/'+args.arch+'_'+ args.dataset + '/'
+    save_path = args.root + '/inverse_loss'+ '/'+args.arch+'_'+ args.dataset + '/'
     adaptation_steps = args.adaptation_steps
     now = datetime.now()
     save_path = save_path + '/' + f'{now.month}_{now.day}_{now.hour}_{now.minute}_{now.second}/'
     os.makedirs(save_path, exist_ok=True)
     wandb.log({'save path': save_path})
     save_args_to_file(args, save_path+"args.json")
-    trainset_ori, testset_ori = get_dataset('ImageNet', '../../../datasets/', subset='imagenette', args=args)
+    trainset_ori, testset_ori = get_dataset("PILE", '../../../datasets', args=args) #get_dataset('ImageNet', '../../../datasets/', subset='imagenette', args=args)
     original_trainloader = DataLoader(trainset_ori, batch_size=args.bs, shuffle=True, num_workers=0)
     original_testloader = DataLoader(testset_ori, batch_size=args.bs, shuffle=False, num_workers=0)
     trainset_tar, testset_tar = get_dataset(args.dataset, '../../../datasets', args=args)
@@ -264,12 +267,19 @@ def main(
             except StopIteration:
                 original_iter = iter(original_trainloader)
                 batch = next(original_iter)
-            inputs, targets = batch
-            inputs, targets = inputs.cuda(), targets.cuda()       
-            # print(inputs.shape)
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
             natural_optimizer.zero_grad()
-            outputs = model(inputs)
-            loss = criterion(outputs, targets) 
+            outputs = model(input_ids=inputs, attention_mask=attention_mask, output_hidden_states=True)
+
+            last_hidden_state = outputs.hidden_states[-1]
+            with torch.no_grad():
+                mask = attention_mask == 1
+                next_token_indexes = (mask.cumsum(dim=1) * mask).argmax(dim=1)
+                next_token_indexes[next_token_indexes == 0] = -1
+            last_token_hidden_state = last_hidden_state[range(last_hidden_state.shape[0]), next_token_indexes]
+            logits = lm_model.lm_head(last_token_hidden_state)
+
+            loss = criterion(logits, targets)
             loss.backward()
             avg_gradients = check_gradients(model)
             # print('check gradients!!!!!!!!!')
diff --git a/classification/kl_uniform_loss.py b/classification/kl_uniform_loss.py
index 059ee86..a3ec476 100644
--- a/classification/kl_uniform_loss.py
+++ b/classification/kl_uniform_loss.py
@@ -21,7 +21,7 @@ def args_parser():
     parser.add_argument('--test_iterval', default=10, type=int)
     parser.add_argument('--arch', default='', type=str)
     parser.add_argument('--gpus', default='0,1', type=str)
-    parser.add_argument('--dataset', default='CIFAR10', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC'])
+    parser.add_argument('--dataset', default='CIFAR10', type=str, choices=['CIFAR10', 'MNIST', 'SVHN', 'STL', 'CINIC', 'IMDB'])
     parser.add_argument('--finetune_epochs', default=1, type=int)
     parser.add_argument('--truly_finetune_epochs', default=20, type=int)
     parser.add_argument('--finetune_lr', default=0.0001, type=float)
@@ -133,9 +133,9 @@ def partial_fast_adapt_multibatch(batches, learner, loss, shots, ways, device):
 
 def test_finetune(model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
@@ -153,9 +153,9 @@ def test_finetune(model, trainset, testset, epochs, lr):
 
 def test_finetune_final(mode, model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model)
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
-    optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
+    trainloader = DataLoader(trainset, batch_size=8, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=8, shuffle=False, num_workers=4,drop_last=True)
+    optimizer = optim.SGD(model.parameters(), lr=lr, weight_decay=0.01)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
     model.train()
diff --git a/classification/model.py b/classification/model.py
index 500d18f..091c91f 100644
--- a/classification/model.py
+++ b/classification/model.py
@@ -3,8 +3,12 @@
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.model_zoo as model_zoo
+from scipy.special import softmax
+from torch.nn.modules.module import T
 from torchvision.models.resnet import ResNet
 from torchvision.models.resnet import BasicBlock, Bottleneck
+from transformers import GPT2LMHeadModel, GPT2ForSequenceClassification, GPT2Tokenizer
+
 
 # 定义ResNet-18结构
 def resnet18(pretrained=False, **kwargs):
@@ -136,4 +140,14 @@ def vgg19(pretrained=False, **kwargs):
     'vgg13_bn': 'https://download.pytorch.org/models/vgg13_bn-abd245e5.pth',
     'vgg16_bn': 'https://download.pytorch.org/models/vgg16_bn-6c64b313.pth',
     'vgg19_bn': 'https://download.pytorch.org/models/vgg19_bn-c79401a0.pth',
-}
\ No newline at end of file
+}
+
+def gpt2(pretrained=False, **kwargs):
+    model = GPT2ForSequenceClassification.from_pretrained("gpt2", num_labels=2, **kwargs)
+    model.config.num_labels = 2
+
+    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token  # Use eos_token as padding token
+
+    model.config.pad_token_id = tokenizer.pad_token_id
+    return model
diff --git a/classification/utils.py b/classification/utils.py
index 9a0048c..ac45ce5 100644
--- a/classification/utils.py
+++ b/classification/utils.py
@@ -3,6 +3,9 @@
 
 import time
 import numpy as np
+from datasets import load_dataset
+from transformers import GPT2LMHeadModel, GPT2ForSequenceClassification
+from datasets import Dataset as HGDataset
 from torch.nn import CrossEntropyLoss
 import torch
 from torch.utils.data import DataLoader, Dataset
@@ -24,7 +27,6 @@
 from torch.utils.data import DataLoader, Dataset
 import torch.nn.init as init
 import csv
-from lib import VGG, make_layers, cfg
 from PIL import Image
 from typing import (
     Generic,
@@ -38,6 +40,9 @@
     Union,
     Dict
 )
+
+from transformers import GPT2Tokenizer
+
 dataTransform = transforms.Compose([
     transforms.ToPILImage(),
     transforms.Resize((64, 64)),
@@ -372,8 +377,57 @@ def get_dataset(dataset, data_path, subset="imagenette", args=None):
         trainset = stl_Dataset([list_img_train, list_label_train])
         testset = stl_Dataset([list_img_test, list_label_test])
 
+    elif dataset.upper() == 'IMDB':
+        dataset = load_dataset("imdb").shuffle(seed=42)
+        train_dataset, test_dataset = dataset["train"].select(range(25000)), dataset["test"].select(range(1000))
+
+        # tokenization
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
+
+        preprocess_function = lambda examples: tokenizer(examples["text"], truncation=True,  padding="max_length",  max_length=512)
+
+        trainset = train_dataset.map(preprocess_function, batched=True)
+        testset = test_dataset.map(preprocess_function, batched=True)
+
+    elif dataset.lower() == 'pile':
+        # Load the dataset from Hugging Face
+        dataset = load_dataset("EleutherAI/the_pile_deduplicated", streaming=True)
+        small_dataset = dataset['train'].shuffle(seed=42, buffer_size=50000)
+        test_dataset = HGDataset.from_list(list(small_dataset.take(1000)))
+        train_dataset = HGDataset.from_list(list(small_dataset.skip(1000).take(10000)))
+
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token
+
+        def preprocess_function(examples):
+            tokenized_inputs = tokenizer(examples["text"], truncation=True, max_length=512)
+
+            input_ids = [tokens[:-1] for tokens in tokenized_inputs['input_ids']]  # Remove the last token
+            labels = [tokens[-1] for tokens in tokenized_inputs['input_ids']]  # Get the last token as label
+
+            # Pad the input_ids to max_length
+            padded_inputs = tokenizer.pad(
+                {"input_ids": input_ids},
+                padding="max_length",
+                max_length=512,  # One token less to account for truncation
+                return_tensors="pt"  # Return PyTorch tensors
+            )
+
+            attention_mask = (padded_inputs["input_ids"] != tokenizer.pad_token_id).long()
+
+            return {
+                "input_ids": padded_inputs["input_ids"],
+                "attention_mask": attention_mask,
+                "label": labels
+            }
+
+        trainset = train_dataset.map(preprocess_function, batched=True)
+        testset = test_dataset.map(preprocess_function, batched=True) 
+
     else:
         exit('unknown dataset: %s'%dataset)
+
     return trainset, testset
 
 def process(checkpoint):
@@ -393,110 +447,110 @@ def get_default_convnet_setting():
 
 
 
-def get_network(model, channel, num_classes, im_size=(32, 32), dist=True):
-    torch.random.manual_seed(int(time.time() * 1000) % 100000)
-    print(f"----------------Using {model} Model----------------")
-    net_width, net_depth, net_act, net_norm, net_pooling = get_default_convnet_setting()
-
-    if model == 'MLP':
-        net = MLP(channel=channel, num_classes=num_classes)
-    elif model == 'ConvNet':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'LeNet':
-        net = LeNet(channel=channel, num_classes=num_classes)
-    elif model == 'AlexNet':
-        net = AlexNet(channel=channel, num_classes=num_classes)
-    elif model == 'VGG11':
-        net = VGG11( channel=channel, num_classes=num_classes)
-    elif model == 'VGG11BN':
-        net = VGG11BN(channel=channel, num_classes=num_classes)
-    elif model == 'ResNet18':
-        net = ResNet18(channel=channel, num_classes=num_classes)
-    elif model == 'ResNet18BN_AP':
-        net = ResNet18BN_AP(channel=channel, num_classes=num_classes)
-    elif model == 'ResNet18_AP':
-        net = ResNet18_AP(channel=channel, num_classes=num_classes)
-
-    elif model == 'ConvNetD1':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=1, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD2':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=2, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD3':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=3, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD4':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=4, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD5':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=5, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD6':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=6, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD7':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=7, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-    elif model == 'ConvNetD8':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=8, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
-
-
-    elif model == 'ConvNetW32':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=32, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetW64':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=64, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetW128':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=128, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetW256':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=256, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetW512':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=512, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetW1024':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=1024, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
-
-    elif model == "ConvNetKIP":
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=1024, net_depth=net_depth, net_act=net_act,
-                      net_norm="none", net_pooling=net_pooling)
-
-    elif model == 'ConvNetAS':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='sigmoid', net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetAR':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='relu', net_norm=net_norm, net_pooling=net_pooling)
-    elif model == 'ConvNetAL':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='leakyrelu', net_norm=net_norm, net_pooling=net_pooling)
-
-    elif model == 'ConvNetNN':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='none', net_pooling=net_pooling)
-    elif model == 'ConvNetBN':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='batchnorm', net_pooling=net_pooling)
-    elif model == 'ConvNetLN':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='layernorm', net_pooling=net_pooling)
-    elif model == 'ConvNetIN':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='instancenorm', net_pooling=net_pooling)
-    elif model == 'ConvNetGN':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='groupnorm', net_pooling=net_pooling)
-
-    elif model == 'ConvNetNP':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='none')
-    elif model == 'ConvNetMP':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='maxpooling')
-    elif model == 'ConvNetAP':
-        net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='avgpooling')
-
-    ###
-    elif model == 'WideRes':
-        net = WideResNet(channel=channel, num_classes=num_classes)
-
-
-    else:
-        net = None
-        exit('DC error: unknown model')
-
-    if dist:
-        gpu_num = torch.cuda.device_count()
-        if gpu_num>0:
-            device = 'cuda'
-            if gpu_num>1:
-                net = nn.DataParallel(net)
-        else:
-            device = 'cpu'
-        net = net.to(device)
-
-    return net
+# def get_network(model, channel, num_classes, im_size=(32, 32), dist=True):
+#     torch.random.manual_seed(int(time.time() * 1000) % 100000)
+#     print(f"----------------Using {model} Model----------------")
+#     net_width, net_depth, net_act, net_norm, net_pooling = get_default_convnet_setting()
+#
+#     if model == 'MLP':
+#         net = MLP(channel=channel, num_classes=num_classes)
+#     elif model == 'ConvNet':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'LeNet':
+#         net = LeNet(channel=channel, num_classes=num_classes)
+#     elif model == 'AlexNet':
+#         net = AlexNet(channel=channel, num_classes=num_classes)
+#     elif model == 'VGG11':
+#         net = VGG11( channel=channel, num_classes=num_classes)
+#     elif model == 'VGG11BN':
+#         net = VGG11BN(channel=channel, num_classes=num_classes)
+#     elif model == 'ResNet18':
+#         net = ResNet18(channel=channel, num_classes=num_classes)
+#     elif model == 'ResNet18BN_AP':
+#         net = ResNet18BN_AP(channel=channel, num_classes=num_classes)
+#     elif model == 'ResNet18_AP':
+#         net = ResNet18_AP(channel=channel, num_classes=num_classes)
+#
+#     elif model == 'ConvNetD1':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=1, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD2':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=2, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD3':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=3, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD4':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=4, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD5':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=5, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD6':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=6, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD7':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=7, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#     elif model == 'ConvNetD8':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=8, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling, im_size=im_size)
+#
+#
+#     elif model == 'ConvNetW32':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=32, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetW64':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=64, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetW128':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=128, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetW256':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=256, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetW512':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=512, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetW1024':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=1024, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling=net_pooling)
+#
+#     elif model == "ConvNetKIP":
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=1024, net_depth=net_depth, net_act=net_act,
+#                       net_norm="none", net_pooling=net_pooling)
+#
+#     elif model == 'ConvNetAS':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='sigmoid', net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetAR':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='relu', net_norm=net_norm, net_pooling=net_pooling)
+#     elif model == 'ConvNetAL':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act='leakyrelu', net_norm=net_norm, net_pooling=net_pooling)
+#
+#     elif model == 'ConvNetNN':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='none', net_pooling=net_pooling)
+#     elif model == 'ConvNetBN':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='batchnorm', net_pooling=net_pooling)
+#     elif model == 'ConvNetLN':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='layernorm', net_pooling=net_pooling)
+#     elif model == 'ConvNetIN':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='instancenorm', net_pooling=net_pooling)
+#     elif model == 'ConvNetGN':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm='groupnorm', net_pooling=net_pooling)
+#
+#     elif model == 'ConvNetNP':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='none')
+#     elif model == 'ConvNetMP':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='maxpooling')
+#     elif model == 'ConvNetAP':
+#         net = ConvNet(channel=channel, num_classes=num_classes, net_width=net_width, net_depth=net_depth, net_act=net_act, net_norm=net_norm, net_pooling='avgpooling')
+#
+#     ###
+#     elif model == 'WideRes':
+#         net = WideResNet(channel=channel, num_classes=num_classes)
+#
+#
+#     else:
+#         net = None
+#         exit('DC error: unknown model')
+#
+#     if dist:
+#         gpu_num = torch.cuda.device_count()
+#         if gpu_num>0:
+#             device = 'cuda'
+#             if gpu_num>1:
+#                 net = nn.DataParallel(net)
+#         else:
+#             device = 'cpu'
+#         net = net.to(device)
+#
+#     return net
 
 
 
@@ -590,6 +644,7 @@ def resume(resume_path):
     return model
 
 def initialize(args,model): #因为maml会多套一层 所以test_finetune里面的另写一个
+    print(model)
     if args.arch == 'res50':
         last_layer = model.module.module.fc
     elif args.arch == 'caformer': 
@@ -599,7 +654,9 @@ def initialize(args,model): #因为maml会多套一层 所以test_finetune里面
     elif args.arch == 'res34':
         last_layer = model.module.module.fc
     elif args.arch == 'vgg':
-        last_layer == model.module.module.fc
+        last_layer = model.module.module.fc
+    elif args.arch == "gpt2":
+        last_layer = model.module.module.score
     init.xavier_uniform_(last_layer.weight)
     if last_layer.bias is not None:
         init.zeros_(last_layer.bias)
@@ -681,7 +738,7 @@ def get_pretrained_model(args, partial_finetuned=False):
             for param in model.fc.parameters():
                 param.requires_grad = True
         return model.cuda()
-    
+
     elif args.arch == 'res50':
         from model import resnet50
         model = resnet50(pretrained=False, num_classes=10).cuda()
@@ -692,11 +749,46 @@ def get_pretrained_model(args, partial_finetuned=False):
             for param in model.fc.parameters():
                 param.requires_grad = True
         return model.cuda()
+
+    elif args.arch == 'gpt2':
+        from model import gpt2
+        model = gpt2(pretrained=False).cuda()
+        for param in model.parameters():
+            param.requires_grad = False
+        for param in model.score.parameters():
+            param.requires_grad = True
+        return model.cuda()
+
+    elif args.arch == 'gpt2-zeroshot':
+        from model import gpt2_zeroshot
+        model = gpt2_zeroshot(pretrained=False).cuda()
+        # if partial_finetuned:
+        #     for param in model.parameters():
+        #         param.requires_grad = False
+        #     for param in model.fc.parameters():
+        #         param.requires_grad = True
+        return model.cuda()
+
     else:
         assert(0)
 
 def get_finetuned_model(args, our_path, partial_finetuned=False):
-   
+    if args.arch == 'gpt2':
+        model = GPT2ForSequenceClassification.from_pretrained("gpt2", num_labels=2)
+        model.config.num_labels = 2
+        tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+        tokenizer.pad_token = tokenizer.eos_token  # Use eos_token as padding token
+        model.config.pad_token_id = tokenizer.pad_token_id
+
+        state_dict = process(torch.load(our_path)['model'])
+        model.load_state_dict(state_dict)
+        if partial_finetuned:
+            for param in model.parameters():
+                param.requires_grad = False
+            for param in model.head.fc.fc2.parameters():
+                param.requires_grad = True
+        return model.cuda()
+
     if args.arch == 'caformer':
         model = timm.create_model("caformer_m36", pretrained=False)
         classifier = nn.Linear(2304, 10)
@@ -963,13 +1055,12 @@ def test(model, original_testloader, device):
     criterion = nn.CrossEntropyLoss(reduction='sum')
     model.eval()
     with torch.no_grad():
-        for batch_idx, (inputs, targets) in enumerate(original_testloader):
-                
-            inputs, targets = inputs.to(device), targets.to(device)
-            outputs = model(inputs)
-            loss = criterion(outputs, targets) 
+        for batch_idx, batch in enumerate(original_testloader):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()
+            outputs = model(inputs, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, targets)
             test_loss += loss.item()
-            _, predicted = outputs.max(1)
+            _, predicted = outputs.logits.max(1)
             # if batch_idx == 0:
             #     print(f'output is {outputs}') # check model whether NaN
             total += targets.size(0)
@@ -983,15 +1074,23 @@ def test_original(model, original_testloader, device):
     correct = 0
     total = 0
     criterion = nn.CrossEntropyLoss(reduction='sum')
+    lm_model = GPT2LMHeadModel.from_pretrained("gpt2").cuda()
     model.eval()
     with torch.no_grad():
-        for batch_idx, (inputs, targets) in enumerate(original_testloader):
-            inputs, targets = inputs.to(device), targets.to(device)
-            # print(inputs.shape)
-            outputs = model(inputs)
-            loss = criterion(outputs, targets)
+        for batch_idx, batch in enumerate(original_testloader):
+            inputs, targets, attention_mask = torch.stack(batch["input_ids"], dim=1).cuda(), batch["label"].cuda(), torch.stack(batch["attention_mask"], dim=1).cuda()          
+
+            outputs = model(input_ids=inputs, attention_mask=attention_mask, output_hidden_states=True)
+            last_hidden_state = outputs.hidden_states[-1]
+            mask = attention_mask == 1
+            next_token_indexes = (mask.cumsum(dim=1) * mask).argmax(dim=1)
+            next_token_indexes[next_token_indexes == 0] = -1
+            last_token_hidden_state = last_hidden_state[range(last_hidden_state.shape[0]), next_token_indexes]
+            logits = lm_model.lm_head(last_token_hidden_state)
+
+            loss = criterion(logits, targets)
             test_loss += loss.item()
-            _, predicted = outputs.max(1)
+            _, predicted = logits.max(1)
             total += targets.size(0)
             correct += predicted.eq(targets).sum().item()
     acc = 100.*correct/total
@@ -1001,8 +1100,8 @@ def test_original(model, original_testloader, device):
 
 def test_finetune(model, trainset, testset, epochs, lr):
     model = nn.DataParallel(model,device_ids=[0,1])
-    trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4,drop_last=True)
-    testloader = DataLoader(testset, batch_size=256, shuffle=False, num_workers=4,drop_last=True)
+    trainloader = DataLoader(trainset, batch_size=16, shuffle=True, num_workers=4,drop_last=True)
+    testloader = DataLoader(testset, batch_size=16, shuffle=False, num_workers=4,drop_last=True)
     optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
     criterion = nn.CrossEntropyLoss()
     scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
@@ -1075,8 +1174,8 @@ def set_seed(seed):
             testset = DatasetSplit(test_dataset_all, np.squeeze(np.argwhere(np.isin(test_dataset_all.targets, config.img_net_classes))), config.img_net_classes)
             train_dataset_all =datasets.ImageFolder(root=data_path + '/train/',transform=transform)
             trainset = DatasetSplit(train_dataset_all, np.squeeze(np.argwhere(np.isin(train_dataset_all.targets, config.img_net_classes))), config.img_net_classes)   
-            trainloader = DataLoader(trainset, batch_size=256, shuffle=True, num_workers=4)
-            testloader = DataLoader(testset, batch_size=256, shuffle=True, num_workers=4)
+            trainloader = DataLoader(trainset, batch_size=16, shuffle=True, num_workers=4)
+            testloader = DataLoader(testset, batch_size=16, shuffle=True, num_workers=4)
             image_train = []
             image_test = []
             target_train = []
diff --git a/environment.yml b/environment.yml
index a2ea879..502057e 100644
--- a/environment.yml
+++ b/environment.yml
@@ -51,6 +51,8 @@ dependencies:
   - xz=5.4.2=h5eee18b_0
   - zeromq=4.3.4=h2531618_0
   - pip:
+      - transformers==4.44.2
+      - datasets==2.21.0
       - accelerate==0.23.0
       - aiohttp==3.8.4
       - aiosignal==1.3.1
@@ -113,7 +115,6 @@ dependencies:
       - h11==0.14.0
       - holoviews==1.16.0
       - httplib2==0.20.4
-      - huggingface-hub==0.15.1
       - idna==3.4
       - imageio==2.28.1
       - importlib-metadata==6.6.0
@@ -122,7 +123,6 @@ dependencies:
       - itsdangerous==2.1.2
       - jinja2==3.1.2
       - joblib==1.2.0
-      - jsonargparse==4.21.1
       - kiwisolver==1.4.4
       - kornia==0.7.0
       - lazy-loader==0.2
@@ -193,11 +193,10 @@ dependencies:
       - qdldl==0.1.7.post0
       - qpth==0.0.16
       - readchar==4.0.5
-      - requests==2.28.1
+      - requests==2.32.2
       - retry-decorator==1.1.1
       - rich==13.3.5
       - rsa==4.7.2
-      - safetensors==0.3.1
       - scikit-image==0.20.0
       - scikit-learn==1.2.2
       - scipy==1.9.1
@@ -220,14 +219,10 @@ dependencies:
       - tgt==1.4.4
       - threadpoolctl==3.1.0
       - tifffile==2023.4.12
-      - timm==0.9.0
+      - timm==0.8.19.dev0
       - toolz==0.12.0
-      - torch==2.0.1+cu118
-      - torchaudio==2.0.2+cu118
-      - torchmetrics==0.10.3
-      - torchvision==0.15.2+cu118
       - tornado==6.3.1
-      - tqdm==4.65.0
+      - tqdm==4.66.3
       - traitlets==5.9.0
       - triton==2.0.0
       - typeshed-client==2.3.0
diff --git a/local_environment.yml b/local_environment.yml
new file mode 100644
index 0000000..e2a4ea8
--- /dev/null
+++ b/local_environment.yml
@@ -0,0 +1,252 @@
+name: sophon
+dependencies:
+  - cudatoolkit=11.8
+  - _libgcc_mutex=0.1
+  - asttokens=2.0.5
+  - backcall=0.2.0
+  - bzip2=1.0.8
+  - ca-certificates=2023.05.30
+  - comm=0.1.2
+  - debugpy=1.5.1
+  - decorator=5.1.1
+  - executing=0.8.3
+  - importlib_metadata=6.0.0
+  - ipykernel=6.19.2
+  - ipython=8.12.0
+  - jedi=0.18.1
+  - jupyter_client=8.1.0
+  - jupyter_core=5.3.0
+  - libffi=3.4.4
+  - libsodium=1.0.18
+  - matplotlib-inline=0.1.6
+  - nest-asyncio=1.5.6
+  - openssl=3.0.9
+  - parso=0.8.3
+  - pexpect=4.8.0
+  - pickleshare=0.7.5
+  - pip=23.1.2
+  - platformdirs=2.5.2
+  - prompt-toolkit=3.0.36
+  - ptyprocess=0.7.0
+  - pure_eval=0.2.2
+  - pygments=2.15.1
+  - python=3.9.16
+  - python-dateutil=2.8.2
+  - pyzmq=25.1.0
+#  - readline=8.2
+  - six=1.16.0
+  - stack_data=0.2.0
+  - tk=8.6.12
+  - typing_extensions=4.6.3
+  - wheel=0.38.4
+  - xz=5.4.2
+  - zeromq=4.3.4
+  - pip:
+      - datasets==2.21.0
+      - transformers==4.44.2
+      - accelerate==0.23.0
+      - aiohttp==3.8.4
+      - aiosignal==1.3.1
+      - alembic==1.12.0
+      - anyio==3.6.2
+      - appdirs==1.4.4
+      - argcomplete==3.1.2
+      - arrow==1.2.3
+      - async-timeout==4.0.2
+      - attrs==23.1.0
+      - audioread==3.0.0
+      - beautifulsoup4==4.12.2
+      - bleach==6.0.0
+      - blessed==1.20.0
+      - bokeh==2.4.3
+      - boto==2.49.0
+      - cachetools==5.3.1
+      - certifi==2022.12.7
+      - cffi==1.15.1
+      - charset-normalizer==2.1.1
+      - clarabel==0.6.0
+      - click==8.1.3
+      - cloudpickle==2.2.1
+      - cmake==3.25.0
+      - colorcet==3.0.1
+      - colorlog==6.7.0
+      - contourpy==1.0.7
+      - crcmod==1.7
+      - croniter==1.3.14
+      - cryptography==41.0.4
+      - cvxpy==1.4.1
+      - cycler==0.11.0
+      - cython==3.0.2
+      - dask==2023.4.1
+      - datashader==0.14.4
+      - datashape==0.5.2
+      - dateutils==0.6.12
+      - deepdiff==6.3.0
+      - detectors==0.1.10
+      - docker-pycreds==0.4.0
+      - docstring-parser==0.15
+      - ecos==2.0.12
+      - faiss-cpu==1.7.4
+      - fastapi==0.88.0
+      - fasteners==0.19
+      - filelock==3.9.0
+      - fonttools==4.39.4
+      - frozenlist==1.3.3
+      - fsspec==2023.5.0
+      - gcs-oauth2-boto-plugin==3.0
+      - gitdb==4.0.10
+      - gitpython==3.1.31
+      - google-apitools==0.5.32
+      - google-auth==2.23.3
+      - google-reauth==0.1.1
+      - greenlet==3.0.0
+      - gsutil==5.26
+      - gym==0.26.2
+      - gym-notices==0.0.8
+      - h11==0.14.0
+      - holoviews==1.16.0
+      - httplib2==0.20.4
+      - huggingface-hub==0.15.1
+      - idna==3.4
+      - imageio==2.28.1
+      - importlib-metadata==6.6.0
+      - importlib-resources==5.12.0
+      - inquirer==3.1.3
+      - itsdangerous==2.1.2
+      - jinja2==3.1.2
+      - joblib==1.2.0
+#      - jsonargparse==4.21.1
+      - kiwisolver==1.4.4
+      - kornia==0.7.0
+      - lazy-loader==0.2
+      - learn2learn==0.2.0
+      - librosa==0.10.1
+      - lightning==2.0.2
+      - lightning-cloud==0.5.36
+      - lightning-flash==0.8.1.post0
+      - lightning-utilities==0.8.0
+      - lit==15.0.7
+      - littleutils==0.2.2
+      - llvmlite==0.40.0
+      - locket==1.0.0
+      - mako==1.2.4
+      - markdown==3.4.3
+      - markdown-it-py==2.2.0
+      - markupsafe==2.1.2
+      - matplotlib==3.7.1
+      - mdurl==0.1.2
+      - monotonic==1.6
+      - mpmath==1.2.1
+      - msgpack==1.0.5
+      - multidict==6.0.4
+      - multipledispatch==0.6.0
+      - munkres==1.1.4
+      - networkx==3.0
+      - numba==0.57.0
+      - numpy==1.23.5
+      - oauth2client==4.1.3
+      - ogb==1.3.6
+      - opencv-python==4.7.0.72
+      - optuna==3.4.0
+      - ordered-set==4.1.0
+      - osqp==0.6.3
+      - outdated==0.2.2
+      - packaging==23.1
+      - pandas==2.0.1
+      - panel==0.14.4
+      - param==1.13.0
+      - partd==1.4.0
+      - pathtools==0.1.2
+      - pillow==9.3.0
+      - plotly==5.13.1
+      - pooch==1.7.0
+      - protobuf==3.20.1
+      - psutil==5.9.5
+      - pyasn1==0.5.0
+      - pyasn1-modules==0.3.0
+      - pybind11==2.11.1
+      - pycparser==2.21
+      - pyct==0.5.0
+      - pydantic==1.10.7
+      - pydeprecate==0.3.2
+      - pyjwt==2.7.0
+      - pynndescent==0.5.10
+      - pyopenssl==23.2.0
+      - pyparsing==3.0.9
+      - python-editor==1.0.4
+      - python-multipart==0.0.6
+      - pytorch-fid==0.3.0
+      - pytorch-lightning==1.9.0
+      - pytz==2023.3
+      - pyu2f==0.1.5
+      - pyviz-comms==2.2.1
+      - pywavelets==1.4.1
+      - pyworld==0.3.4
+      - pyyaml==6.0
+      - qdldl==0.1.7.post0
+      - qpth==0.0.16
+      - readchar==4.0.5
+      - requests==2.28.1
+      - retry-decorator==1.1.1
+      - rich==13.3.5
+      - rsa==4.7.2
+      - safetensors==0.3.1
+      - scikit-image==0.20.0
+      - scikit-learn==1.2.2
+      - scipy==1.9.1
+      - scs==3.2.3
+      - seaborn==0.12.2
+      - sentry-sdk==1.23.0
+      - setproctitle==1.3.2
+      - setuptools==59.5.0
+      - smmap==5.0.0
+      - sniffio==1.3.0
+      - soundfile==0.12.1
+      - soupsieve==2.4.1
+      - soxr==0.3.6
+      - sqlalchemy==2.0.22
+      - starlette==0.22.0
+      - starsessions==1.3.0
+      - sympy==1.11.1
+      - tenacity==8.2.2
+      - termcolor==2.3.0
+      - tgt==1.4.4
+      - threadpoolctl==3.1.0
+      - tifffile==2023.4.12
+#      - timm==0.9.0
+      - toolz==0.12.0
+      - torch==2.0.1
+      - torchaudio==2.0.2
+      - torchmetrics==0.10.3
+      - torchvision==0.15.2
+      - tornado==6.3.1
+      - tqdm==4.65.0
+      - traitlets==5.9.0
+#      - triton==2.0.0
+      - typeshed-client==2.3.0
+      - typing-extensions==4.4.0
+      - tzdata==2023.3
+      - umap-learn==0.5.3
+      - urllib3==1.26.13
+      - uvicorn==0.22.0
+      - wandb==0.15.2
+      - wcwidth==0.2.6
+      - webencodings==0.5.1
+      - websocket-client==1.5.1
+      - websockets==11.0.3
+      - wilds==2.0.0
+      - xarray==2023.4.2
+      - yarl==1.9.2
+      - zipp==3.15.0
+#      - ncurses
+#      - libnsl==2.0.0
+#      - libsqlite==3.41.2
+#      - _openmp_mutex==4.5
+#      - ld_impl_linux-64==2.38
+#      - libzlib==1.2.13
+#      - libgomp==12.2.0
+#      - libstdcxx-ng==11.2.0
+#      - libgcc-ng==12.2.0
+#      - libuuid==2.38.1
+#      - readline==8.2
+prefix: /opt/anaconda3/envs/diffuser