Commit
·
b8dafec
1
Parent(s):
7fdab2a
updated readme
Browse files
README.md
CHANGED
|
@@ -122,11 +122,15 @@ checkpoint_path = hf_hub_download(
|
|
| 122 |
)
|
| 123 |
|
| 124 |
# 2️⃣ Load the model weights safely
|
| 125 |
-
state_dict = load_file(checkpoint_path)
|
| 126 |
model = ImprovedEfficientViT()
|
| 127 |
model.load_state_dict(state_dict)
|
| 128 |
model.eval()
|
| 129 |
|
|
|
|
|
|
|
|
|
|
|
|
|
| 130 |
# 3️⃣ Run inference on a video
|
| 131 |
video_path = "sample_video.mp4"
|
| 132 |
result = predict_vedio(video_path, model)
|
|
|
|
| 122 |
)
|
| 123 |
|
| 124 |
# 2️⃣ Load the model weights safely
|
| 125 |
+
state_dict = load_file(checkpoint_path, device="cpu")
|
| 126 |
model = ImprovedEfficientViT()
|
| 127 |
model.load_state_dict(state_dict)
|
| 128 |
model.eval()
|
| 129 |
|
| 130 |
+
# 4️⃣ Move to GPU if available
|
| 131 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 132 |
+
model.to(device)
|
| 133 |
+
|
| 134 |
# 3️⃣ Run inference on a video
|
| 135 |
video_path = "sample_video.mp4"
|
| 136 |
result = predict_vedio(video_path, model)
|