Solved my issue. Needed this:
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", torch_dtype=torch.float16, low_cpu_mem_usage=True).to("cuda",  torch_dtype=torch.float16)

Also need to move the input ids to the GPU:
input_ids = tokenizer(context, return_tensors="pt").input_ids.to("cuda")