Each element is 200 floats so 800 bytes.
Here is the brief version of my code (with only minimal memory leak patches):
# pad_len := max number of tokens (input sequences padded to this length)
def train_loop(model, input_output_data, epochs=5, batch_size=32, pad_len=200):
    input, output = input_output_data
    n_examples = len(input)
    
    model.train() # turn on training
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    n_batches = int(n_examples/batch_size+0.99999)
    for i in range(epochs):
        print(f'epoch: {i+1}/{epochs}')
        for j in range(n_batches):
            optimizer.zero_grad()
            torch.cuda.empty_cache()
            try:
                # they must be padded to the same size for batching to work...
                inputs = tokenizer(input[j*batch_size:(j+1)*batch_size], return_tensors='pt',
                                   padding='max_length', truncation=True, max_length=pad_len)
                outputs = tokenizer(output[j*batch_size:(j+1)*batch_size], return_tensors='pt',
                                    padding='max_length', truncation=True, max_length=pad_len)
                loss = model(**inputs, labels=outputs['input_ids'])[0].mean()
            except Exception as e:
                #pdb.set_trace()
                raise
            
            print(f'batch: {j+1}/{n_batches}, loss: {loss}')
            loss.backward() # mean is unnecessary but just for safety
            optimizer.step()