| 113 | } |
| 114 | |
| 115 | func TestDevice(t *testing.T) { |
| 116 | s := NewScope() |
| 117 | matrix := Const(s, [][]float32{{3.0}}) |
| 118 | s = s.WithDevice("/device:GPU:0") |
| 119 | square := MatMul(s.SubScope("square"), matrix, matrix) |
| 120 | s = s.WithDevice("") |
| 121 | cube := MatMul(s.SubScope("cube"), square, matrix) |
| 122 | if got, want := square.Op.Device(), "/device:GPU:0"; got != want { |
| 123 | t.Errorf("Got %q, want %q", got, want) |
| 124 | } |
| 125 | if got, want := cube.Op.Device(), ""; got != want { |
| 126 | t.Errorf("Got %q, want %q", got, want) |
| 127 | } |
| 128 | } |
| 129 | |
| 130 | func TestScopeFinalize(t *testing.T) { |
| 131 | var ( |