# -*- coding: utf-8 -*-
"""
Created on Sun Sep 12 23:31:16 2021

@author: arash
"""

import torch, pickle
from training_net import trainingNetwork, trainingLineNetwork
from data_loader import my_collate, CustomIterableDataset, CustomDataset, CustomDatasetBalanced
import random
import numpy as np

from test_line import test


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

available_gpus = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
print('Number of GPUs: ', len(available_gpus))

#dataset = CustomIterableDataset(data_link='training_processed_LLVM3.p')
#dataset = CustomDataset()
#dataset = CustomDatasetBalanced(data_link='training_processed_LLVM4.p')
#print('dataset')
n_gpus = torch.cuda.device_count()
tot=[]
for _ in range(5):

    Net, loss=trainingLineNetwork()
    res=test(Net)
    tot.append(res)
    pickle.dump(tot, open('res_line_5_runs.p','wb'))

#Net, loss=trainingNetwork(dataset).cuda()
#Net=nn.DataParallel(Net)
#Net.to(device)

#pickle.dump(loss, open('loss.p','wb'))
