Update README.md

README.md

@@ -61,11 +61,7 @@ test_set = model.get_dataset(database_dir=database_dir, test_set_only=True)
 
 # Get an example, add mini-batch dimension and move to device:
 example = test_set[0]
-for k, v in example.items():
-    if isinstance(v, torch.Tensor):
-        example[k] = v.to(device).unsqueeze(0)
-    else:
-        example[k] = [v]
+example = {k: v.to(device).unsqueeze(0) if isinstance(v, torch.Tensor) else [v] for k, v in example.items()}  # Add mini-batch dimension and move to device.
 
 # Convert the patient data in the batch into embeddings:
 inputs_embeds, attention_mask, token_type_ids, position_ids, bos_token_ids = model.prepare_inputs(tokenizer=tokenizer, **example)