-
Notifications
You must be signed in to change notification settings - Fork 4
/
cuda.py
25 lines (20 loc) · 1.06 KB
/
cuda.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#!/usr/bin/env python
# -*- coding: utf-8 -*-
################################################################################
# Copyright (c) 2017. Vincenzo Lomonaco, Karan Desai, Eugenio Culurciello, #
# Davide Maltoni. All rights reserved. #
# See the accompanying LICENSE file for terms. #
# #
# Date: 27-05-2019 #
# Authors: Vincenzo Lomonaco, Karan Desai, Eugenio Culurciello, Davide Maltoni #
# E-mail: vincenzo.lomonaco@unibo.it #
# Website: vincenzolomonaco.com #
################################################################################
import torch
import torch.autograd as autograd
USE_CUDA = torch.cuda.is_available()
def Variable(data, *args, **kwargs):
var = autograd.Variable(data, *args, **kwargs)
if USE_CUDA:
var = var.cuda()
return var