Hi yes
from torch.utils.data import Dataset, DataLoader
import numpy as np
"""Behavior Cloning Dataset"""
class BCDataset(Dataset):
  def __init__(self, df=None, transform=None):
    self.ins_and_act_df = df if df else pd.read_csv('instructions_and_actions.csv')
    self.n_samples = self.ins_and_act_df.shape[0]
    self.transform = transform

  def __getitem__(self, index):
    # Obtain text
    encoding_perceiver = tokenizer(self.ins_and_act_df.iloc[index, 1], padding=True, return_tensors="pt")
    inputs, input_mask = encoding_perceiver.input_ids, encoding_perceiver.attention_mask
    outputs_perceiver = model_perceiver(inputs=inputs, attention_mask=input_mask)
    hidden_states_perceiver = outputs_perceiver.hidden_states[-1][:, 0, :].detach()

    # Obtain image
    img = Image.open('images/{}.png'.format(index))
    encoding_vit = feature_extractor(images=img, return_tensors="pt")
    pixel_values = encoding_vit['pixel_values']
    outputs_vit = model_vit(pixel_values)
    last_hidden_states = outputs_vit.last_hidden_state[:, 0, :].detach()

    # Obtain labels
    y = torch.from_numpy(np.array(self.ins_and_act_df.iloc[index, 2:], dtype=np.int8))

    if self.transform:
        img = self.transform(img)

    x = torch.cat((hidden_states_perceiver, last_hidden_states), axis=1).to(device)
    return torch.nn.functional.normalize(x), y.to(device)

  def __len__(self):
    return self.n_samples

dataset = BCDataset()
train_set, val_set = torch.utils.data.random_split(dataset, [int(dataset.n_samples * 0.8),int(dataset.n_samples * 0.2)])
train_loader = DataLoader(dataset = train_set, batch_size=4, shuffle=True, num_workers=2)
val_loader = DataLoader(dataset = val_set, batch_size=4, shuffle=True, num_workers=2)