Here is the code. I’m using a p3.2xlarge SageMaker notebook instance. If I run the code with a conda_pytorch_p36 kernel I get the error
`

ImportError: Using FP16 with APEX but APEX is not installed, please refer to GitHub - NVIDIA/apex: A PyTorch Extension: Tools for easy mixed precision and distributed training in Pytorch.

`
if I run it with the pytorch_latest I get
> 
> ~/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/trainer.py in compute_loss(self, model, inputs, return_outputs)
>    1774         else:
>    1775             # We don't use .loss here since the model may return tuples instead of ModelOutput.
> -> 1776             loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
>    1777 
>    1778         return (loss, outputs) if return_outputs else loss
> 
> ~/anaconda3/envs/pytorch_latest_p36/lib/python3.6/site-packages/transformers/file_utils.py in __getitem__(self, k)
>    1736         if isinstance(k, str):
>    1737             inner_dict = {k: v for (k, v) in self.items()}
> -> 1738             return inner_dict[k]
>    1739         else:
>    1740             return self.to_tuple()[k]
> 
> KeyError: 'loss'

I believe I’m getting the ‘loss’ key error because I’m doing something wrong with the training arguments. I have been trying to pass different loss functions etc, but no luck so far.
Full code to reproduce on p3.2xlarge below - thanks for the help!

!pip install "sagemaker>=2.48.0" "transformers==4.6.1" "datasets[s3]==1.6.2" --upgrade

# Download some training data

!wget https://github.com/saurabh3949/Text-Classification-Datasets/raw/master/dbpedia_csv.tar.gz
!tar -xzvf dbpedia_csv.tar.gz

import pandas as pd
import json
import os

# Write small train and test files

df = pd.read_csv('dbpedia_csv/train.csv', header = None)

# write as small train input file
with open('train_text.json', 'w') as outfile:
    for desc in df.iloc[:10000, 2]:
        json.dump({"inputs": desc}, outfile)
        outfile.write('\n')
        
with open('test_text.json', 'w') as outfile:
    for desc in df.iloc[10000:15000, 2]:
        json.dump({"inputs": desc}, outfile)
        outfile.write('\n')

!sudo chmod 777 /opt/ml

from transformers import AutoTokenizer, AutoModel
from transformers import BertTokenizer, BertForSequenceClassification
import torch
import torch.nn.functional as F


args = lambda: None

from transformers import (
    AutoModelForSequenceClassification,
    AutoModel,
    Trainer,
    TrainingArguments,
    AutoTokenizer,
    default_data_collator,
)

from datasets import load_dataset
import sys
import os

raw_train_dataset = load_dataset("json", data_files='train_text.json')["train"]
raw_test_dataset = load_dataset("json", data_files='test_text.json')["train"]

args.model_id = 'facebook/bart-large-mnli'

tokenizer = AutoTokenizer.from_pretrained(args.model_id)
model = AutoModel.from_pretrained(args.model_id)


# preprocess function, tokenizes text
def preprocess_function(examples):
    return tokenizer(examples["inputs"], padding = 'max_length', max_length=52, truncation=True)

# preprocess dataset
train_dataset = raw_train_dataset.map(
    preprocess_function,
    batched=True,
)
test_dataset = raw_test_dataset.map(
    preprocess_function,
    batched=True,
)


# define training args
training_args = TrainingArguments(
    output_dir='/opt/ml/model',
    num_train_epochs=1,
    per_device_train_batch_size=32,
    warmup_steps=500,
    fp16=True,
    do_eval = False,
    save_strategy="epoch",
    logging_dir="s3://tmp232/sourcedir.tar.gz/logs",
    learning_rate=float(3e-05)
)

trainer =Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=tokenizer,
    data_collator=default_data_collator
)

# train model
trainer.train()