diff options
Diffstat (limited to 'tensorflow/go/op/scope.go')
-rw-r--r-- | tensorflow/go/op/scope.go | 31 |
1 files changed, 27 insertions, 4 deletions
diff --git a/tensorflow/go/op/scope.go b/tensorflow/go/op/scope.go index 13de4294dc..ac39808d83 100644 --- a/tensorflow/go/op/scope.go +++ b/tensorflow/go/op/scope.go @@ -37,6 +37,7 @@ type Scope struct { namemap map[string]int namespace string controlDependencies []*tf.Operation + device string err *scopeErr } @@ -82,6 +83,7 @@ func (s *Scope) AddOperation(args tf.OpSpec) *tf.Operation { args.Name = s.namespace + "/" + args.Name } args.ControlDependencies = append(args.ControlDependencies, s.controlDependencies...) + args.Device = s.device op, err := s.graph.AddOperation(args) if err != nil { s.UpdateErr(args.Type, err) @@ -98,10 +100,12 @@ func (s *Scope) SubScope(namespace string) *Scope { namespace = s.namespace + "/" + namespace } return &Scope{ - graph: s.graph, - namemap: make(map[string]int), - namespace: namespace, - err: s.err, + graph: s.graph, + namemap: make(map[string]int), + namespace: namespace, + controlDependencies: s.controlDependencies, + device: s.device, + err: s.err, } } @@ -123,6 +127,25 @@ func (s *Scope) WithControlDependencies(ops ...*tf.Operation) *Scope { namemap: s.namemap, namespace: s.namespace, controlDependencies: deps, + device: s.device, + err: s.err, + } +} + +// WithDevice returns a new Scope which will cause all operations added to the +// graph to execute on devices that match the provided device specification. +// +// For example, WithDevice("/device:GPU:0") will cause operations added to +// the graph to execute on GPU #0. +// +// An empty string removes any device restrictions. +func (s *Scope) WithDevice(device string) *Scope { + return &Scope{ + graph: s.graph, + namemap: s.namemap, + namespace: s.namespace, + controlDependencies: s.controlDependencies, + device: device, err: s.err, } } |