Keep pinned-memory

app.py

@@ -45,8 +45,8 @@ def compile_model():
     weights_: dict[str, torch.Tensor] = {}
     for name in weights:
         tensor, _properties = weights.get_weight(name)
-        tensor_ = torch.empty_like(tensor, device='cpu', requires_grad=False)
-        weights_[name] = tensor_.copy_(tensor).share_memory_() # TODO: non_blocking=True?
+        tensor_ = torch.empty_like(tensor, device='cpu', requires_grad=False).pin_memory()
+        weights_[name] = tensor_.copy_(tensor).share_memory_()
     return weights_
 
 weights = compile_model()