Update README.md

README.md

@@ -39,7 +39,6 @@ import requests
 from PIL import Image
 import torch
 
-device = torch.device('cuda')
 image_urls = [
     "https://farm4.staticflickr.com/3395/3428278415_81c3e27f15_z.jpg",
     "http://images.cocodataset.org/val2017/000000039769.jpg"]
@@ -49,13 +48,13 @@ texts = [
 images = [Image.open(requests.get(url, stream=True).raw) for url in image_urls]
 
 processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-large-itm-mlm")
-model = BridgeTowerForContrastiveLearning.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")model.to(device)
+model = BridgeTowerForContrastiveLearning.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")
 
-inputs  = processor(images, texts, padding=True, return_tensors="pt").to(device)
+inputs  = processor(images, texts, padding=True, return_tensors="pt")
 outputs = model(**inputs, labels=torch.ones(2,device=device))
 
-inputs  = processor(images, texts[::-1], padding=True, return_tensors="pt").to(device)
-outputs_swapped = model(**inputs, labels=torch.ones(2,device=device))
+inputs  = processor(images, texts[::-1], padding=True, return_tensors="pt")
+outputs_swapped = model(**inputs, labels=torch.ones(2))
 
 print('Loss', outputs.loss.item())
 print('Loss with swapped images', outputs_swapped.loss.item())