| 64 | torch.save(value, self.path, _use_new_zipfile_serialization=False) |
| 65 | |
| 66 | class TorchSaveJitStream_CUDA(FileSetup): |
| 67 | path = 'saved_stream_model.pt' |
| 68 | |
| 69 | def setup(self): |
| 70 | if not torch.cuda.is_available(): |
| 71 | return |
| 72 | |
| 73 | class Model(torch.nn.Module): |
| 74 | def forward(self): |
| 75 | s = torch.cuda.Stream() |
| 76 | a = torch.rand(3, 4, device="cuda") |
| 77 | b = torch.rand(3, 4, device="cuda") |
| 78 | |
| 79 | with torch.cuda.stream(s): |
| 80 | is_stream_s = torch.cuda.current_stream(s.device_index()).id() == s.id() |
| 81 | c = torch.cat((a, b), 0).to("cuda") |
| 82 | s.synchronize() |
| 83 | return is_stream_s, a, b, c |
| 84 | |
| 85 | model = Model() |
| 86 | |
| 87 | # Script the model and save |
| 88 | script_model = torch.jit.script(model) |
| 89 | torch.jit.save(script_model, self.path) |
| 90 | |
| 91 | |
| 92 | tests = [ |
no outgoing calls
no test coverage detected
searching dependent graphs…