zluda llm/vlm temp fix

modules/zluda_hijacks.py

@@ -1,7 +1,7 @@
 from functools import wraps
 import torch
 import torch._dynamo.device_interface
-from modules import shared, rocm, zluda # pylint: disable=unused-import
+from modules import shared, zluda # pylint: disable=unused-import
 
 
 MEM_BUS_WIDTH = {
@@ -18,6 +18,13 @@ MEM_BUS_WIDTH = {
 }
 
 
+_topk = torch.topk
+def topk(input: torch.Tensor, *args, **kwargs): # pylint: disable=redefined-builtin
+    device = input.device
+    values, indices = _topk(input.cpu(), *args, **kwargs)
+    return torch.return_types.topk((values.to(device), indices.to(device),))
+
+
 class DeviceProperties:
     PROPERTIES_OVERRIDE = {"regs_per_multiprocessor": 65535, "gcnArchName": "UNKNOWN ARCHITECTURE"}
     internal: torch._C._CudaDeviceProperties
@@ -42,6 +49,7 @@ def torch__C__cuda_getCurrentRawStream(device):
 
 
 def do_hijack():
+    torch.topk = topk
     if zluda.default_agent is not None:
         DeviceProperties.PROPERTIES_OVERRIDE["gcnArchName"] = zluda.default_agent.name
     torch.cuda._get_device_properties = torch_cuda__get_device_properties # pylint: disable=protected-access