# -*- coding: utf-8 -*-
"""torch support.
"""

from timeit import default_timer as timer
import torch
import numpy as np
from .i_gemm import IGEMM


class Benchmark(IGEMM):
    def __init__(self, params, remaining_args=None):
        super().__init__(params, remaining_args)

    def run(self):
        M, N, K = self.matrix_size
        dtype = np.float32
        a = torch.tensor(np.random.random((M, N)).astype(dtype))
        b = torch.tensor(np.random.random((N, K)).astype(dtype))
        c = torch.tensor(np.random.random((M, K)).astype(dtype))
        if "nb_gpus" in self.params:
            if self.params["nb_gpus"] > 1:
                raise RuntimeError("Only 1 GPU is supported")
            if self.params["nb_gpus"] == 1:
                torch.cuda.set_device(self.params["gpus"][0])
                a.to("cuda")
                b.to("cuda")
                c.to("cuda")
        nb_epoch = 2
        time_start = timer()
        for _ in range(nb_epoch):
            c = a @ b  # + c
        time_end = timer()
        elapsed_time = (time_end - time_start) / nb_epoch
        self.params["time"] = elapsed_time
        self.params["GFLOP/sec"] = self.params["GFLOP"] / elapsed_time

@@ -751,52 +751,130 @@
     
-torch.cuda.set_device(self.params%5B%22gpus%22%5D%5B0%5D
+device = torch.device(%22cuda%22)%0A                id_gpu = self.params%5B%22gpus%22%5D%5B0%5D%0A                torch.cuda.set_device(id_gpu
 )%0A  
@@ -888,26 +888,30 @@
        a
-.to(%22cuda%22
+ = a.to(device
 )%0A      
@@ -921,26 +921,30 @@
        b
-.to(%22cuda%22
+ = b.to(device
 )%0A      
@@ -958,18 +958,22 @@
    c
-.to(%22cuda%22
+ = c.to(device
 )%0A  

fix torch to gpu