| 112 | |
| 113 | |
| 114 | def safe_state(silent): |
| 115 | old_f = sys.stdout |
| 116 | class F: |
| 117 | def __init__(self, silent): |
| 118 | self.silent = silent |
| 119 | |
| 120 | def write(self, x): |
| 121 | if not self.silent: |
| 122 | if x.endswith("\n"): |
| 123 | old_f.write(x.replace("\n", " [{}]\n".format(str(datetime.now().strftime("%d/%m %H:%M:%S"))))) |
| 124 | else: |
| 125 | old_f.write(x) |
| 126 | |
| 127 | def flush(self): |
| 128 | old_f.flush() |
| 129 | |
| 130 | sys.stdout = F(silent) |
| 131 | |
| 132 | random.seed(0) |
| 133 | np.random.seed(0) |
| 134 | torch.manual_seed(0) |
| 135 | torch.cuda.set_device(torch.device("cuda:0")) |
| 136 | |
| 137 | def knn(x, src, k, transpose=False): |
| 138 | if transpose: |