args = parser.parse_args() torch.manual_seed(args.seed) if torch.cuda.is_available(): if not args.cuda: print("WARNING: You have a CUDA device, so you should probably run with --cuda") input_channels = 2 n_classes = 1 batch_size = args.batch_size seq_length = args.seq_len epochs = args.epochs print(args) print("Producing data...") X_train, Y_train = data_generator(50000, seq_length) X_test, Y_test = data_generator(1000, seq_length) # Note: We use a very simple setting here (assuming all levels have the same # of channels. channel_sizes = [args.nhid]*args.levels kernel_size = args.ksize dropout = args.dropout model = TCN(input_channels, n_classes, channel_sizes, kernel_size=kernel_size, dropout=dropout) if args.cuda: model.cuda() X_train = X_train.cuda() Y_train = Y_train.cuda() X_test = X_test.cuda() Y_test = Y_test.cuda()