Update README.md

README.md

@@ -99,7 +99,7 @@ inputs = {k: v.to(device) for k, v in feature_extractor(img, return_tensors='pt'
 
 # extract patch embeddings
 with torch.no_grad():
-  patch_embeddings = rearrange(model(**inputs).last_hidden_state, 'b d h w -> b (h w) d').cpu()
+  patch_embeddings = model(**inputs).last_hidden_state[0].permute(1, 2, 0).view(7*7, 2048).cpu()
 
 # classify
 pred = logistic_regression.predict(patch_embeddings.sum(dim=0, keepdim=True))