Update app.py
Browse files
app.py
CHANGED
|
@@ -141,6 +141,7 @@ def create_demo():
|
|
| 141 |
|
| 142 |
|
| 143 |
if __name__ == "__main__":
|
|
|
|
| 144 |
model = Mamba(num_layers = 2, d_input = 1024, d_model = 512, num_classes=7, model_name='jina', pooling=None).to(device)
|
| 145 |
checkpoint = torch.load("models/Mamba_jina_checkpoint.pth", map_location=torch.device('cpu'))
|
| 146 |
model.load_state_dict(checkpoint['model_state_dict'])
|
|
|
|
| 141 |
|
| 142 |
|
| 143 |
if __name__ == "__main__":
|
| 144 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 145 |
model = Mamba(num_layers = 2, d_input = 1024, d_model = 512, num_classes=7, model_name='jina', pooling=None).to(device)
|
| 146 |
checkpoint = torch.load("models/Mamba_jina_checkpoint.pth", map_location=torch.device('cpu'))
|
| 147 |
model.load_state_dict(checkpoint['model_state_dict'])
|