cc_staff
284
edits
No edit summary |
(Fixed set_num_threads expects an int, but got str) |
||
Line 183: | Line 183: | ||
args = parser.parse_args() | args = parser.parse_args() | ||
torch.set_num_threads(os.environ['SLURM_CPUS_PER_TASK']) | torch.set_num_threads(int(os.environ['SLURM_CPUS_PER_TASK'])) | ||
class Net(nn.Module): | class Net(nn.Module): | ||