Make some other changes to allow the loading to transfer from CPU to GPU if appropriate.

theano/sandbox/cuda/__init__.py

@@ -2,6 +2,7 @@ import atexit, logging, os, stat, sys
 from theano.compile import optdb
 from theano.gof.cmodule import get_lib_extension
 from theano.configparser import config, AddConfigVar, StrParam
+from theano.tensor.sharedvar import load_shared_variable
 import nvcc_compiler
 
 _logger_name = 'theano.sandbox.cuda'
@@ -279,22 +280,20 @@ def handle_shared_float32(tf):
         theano.compile.shared_constructor(float32_shared_constructor)
         # this is a bit of hackery to make the shared variables load
         # with the proper type.
-        copy_reg.pickle(theano.gof.graph.Apply, reduce_apply, 
-                        load_shared_pickle)
+        copy_reg.pickle(theano.tensor.basic.TensorVariable, 
+                        reduce_tensor_variable, 
+                        load_shared_variable)
     else:
         raise NotImplementedError('removing our handler')
 
-def reduce_apply(apply):
-    if isinstance(apply.op, HostFromGpu) and len(apply.inputs) == 1 and \
-            isinstance(apply.inputs[0], CudaNdarraySharedVariable):
-        return load_shared_pickle, apply.inputs[0].get_value()
+def reduce_tensor_variable(var):
+    if isinstance(var.owner.op, HostFromGpu) and len(var.owner.inputs) == 1 \
+            and isinstance(var.owner.inputs[0], CudaNdarraySharedVariable):
+        return load_shared_variable, (var.owner.inputs[0].get_value(),)
     else:
         # this will make protocol 2 a little bit less efficient
         # but there is no way around it.
-        return apply.__reduce__()
-
-def load_shared_pickle(val):
-    return theano.tensor.as_tensor_variable(theano.shared(val))
+        return var.__reduce__()
 
 if config.device.startswith('gpu'):
     use(device=config.device, force=config.force_device)

theano/tensor/sharedvar.py

@@ -5,9 +5,17 @@ from basic import TensorType, _tensor_py_operators, autocast_int, autocast_float
 from theano.compile import shared_constructor, SharedVariable
 from theano import config
 
+def load_shared_variable(val):
+    return theano.tensor.as_tensor_variable(theano.shared(val))
+
 # _tensor_py_operators is first to have its version of __{gt,ge,lt,le}__
 class TensorSharedVariable(_tensor_py_operators, SharedVariable):
-    pass
+    def __reduce_ex__(self, proto):
+        # This is for loading on the GPU if present.
+        if self.dtype == 'float32':
+            return load_shared_variable, (self.get_value(),)
+        else:
+            return super(TensorSharedVariable, self).__reduce_ex__(proto)
 
 @shared_constructor
 def tensor_constructor(value, name=None, strict=False, allow_downcast=None, borrow=False, broadcastable=None):