Vjepa

[aparnabg]

No description provided.

[gemini-code-assist]

Code Review

This pull request introduces V-JEPA2 model variants, feature extraction, and probe training scripts for clip-level and frame-level action recognition, along with SLURM submission scripts and unit tests. The code review highlights several critical issues that need to be addressed: potential silent feature-to-metadata misalignments and dimension mismatch crashes due to improper exception handling; dimension mismatches where EMBED_DIM is set to 1408 instead of 1024 for ViT-H features; CPU compatibility failures caused by unconditional torch.float16 loading and .half() casting; potential crashes if bbox_map is empty; and portability issues from hardcoded absolute user paths in SLURM scripts.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If a batch fails during feature extraction, catching the exception and continuing without appending to all_features or all_labels will cause the saved features to be shorter than the metadata list in dataset_meta.json. This leads to a silent misalignment (shift) between the features and their corresponding video paths during training and evaluation. It is safer to raise the exception and fail fast.

Suggested change

	except Exception as e:
	print(f" Error at batch {batch_idx}: {e}")
	errors.append(batch_idx)
	except Exception as e:
	print(f" Error at batch {batch_idx}: {e}")
	raise e

[gemini-code-assist]

If a batch fails during feature extraction, catching the exception and continuing without appending to all_features or all_labels will cause the saved features to be shorter than the metadata list in dataset_meta.json. This leads to a silent misalignment (shift) between the features and their corresponding video paths during training and evaluation. It is safer to raise the exception and fail fast.

Suggested change

	except Exception as e:
	print(f" Error at batch {batch_idx}: {e}", flush=True)
	errors.append(batch_idx)
	except Exception as e:
	print(f" Error at batch {batch_idx}: {e}")
	raise e

[gemini-code-assist]

If a batch fails during feature extraction, appending a dummy tensor of shape [B, 1, EMBED_DIM] will cause torch.cat(all_feats, 0) to crash with a RuntimeError because the token dimension (dimension 1) will mismatch with successful batches (which have shape [B, N_tokens, EMBED_DIM]). It is safer to raise the exception and fail fast.

Suggested change

	except Exception as e:
	print(f" batch {i} error: {e}")
	all_feats.append(torch.zeros(clips.shape[0], 1, EMBED_DIM))
	all_labels.append(labels)
	except Exception as e:
	print(f" batch {i} error: {e}")
	raise e

[gemini-code-assist]

The loaded full video features are 1024-dimensional (as indicated by the comments and the ViT-H feature shape), but EMBED_DIM is set to 1408. Initializing the AttentiveProbe with embed_dim=1408 while passing 1024-dimensional features will cause a RuntimeError in nn.MultiheadAttention due to dimension mismatch. Please update EMBED_DIM to 1024.

Suggested change

	EMBED_DIM = 1408
	EMBED_DIM = 1024

[gemini-code-assist]

The loaded full video features are 1024-dimensional (as indicated by the comments and the ViT-H feature shape), but EMBED_DIM is set to 1408. Initializing the HierarchicalProbe with embed_dim=1408 while passing 1024-dimensional features will cause a RuntimeError in nn.MultiheadAttention due to dimension mismatch. Please update EMBED_DIM to 1024.

Suggested change

	EMBED_DIM = 1408
	EMBED_DIM = 1024

[gemini-code-assist]

Loading the model with torch_dtype=torch.float16 unconditionally will cause a RuntimeError on CPU (when GPU is not available) because CPU does not support most float16 operations in PyTorch. Please set the dtype dynamically based on CUDA availability.

        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        self.encoder = AutoModel.from_pretrained(
            MODEL_NAME, torch_dtype=dtype, attn_implementation="sdpa",
        )

[gemini-code-assist]

Manually casting inputs to .half() will crash on CPU when running in float32 mode. It is safer to cast the inputs dynamically to the encoder's dtype.

Suggested change

	inputs = {k: (v.half() if v.is_floating_point() else v) for k, v in inputs.items()}
	inputs = {k: (v.to(self.encoder.dtype) if v.is_floating_point() else v) for k, v in inputs.items()}

[gemini-code-assist]

Loading the model with torch_dtype=torch.float16 unconditionally will cause a RuntimeError on CPU (when GPU is not available) because CPU does not support most float16 operations in PyTorch. Please set the dtype dynamically based on CUDA availability.

Suggested change

	encoder = AutoModel.from_pretrained(
	HF_MODEL_NAME,
	torch_dtype=torch.float16,
	attn_implementation="sdpa",
	)
	dtype = torch.float16 if torch.cuda.is_available() else torch.float32
	encoder = AutoModel.from_pretrained(
	HF_MODEL_NAME,
	torch_dtype=dtype,
	attn_implementation="sdpa",
	)

[gemini-code-assist]

Loading the model with torch_dtype=torch.float16 unconditionally will cause a RuntimeError on CPU (when GPU is not available) because CPU does not support most float16 operations in PyTorch. Please set the dtype dynamically based on CUDA availability.

Suggested change

	processor = AutoVideoProcessor.from_pretrained(HF_MODEL_NAME)
	encoder = AutoModel.from_pretrained(
	HF_MODEL_NAME,
	torch_dtype=torch.float16,
	attn_implementation="sdpa",
	dtype = torch.float16 if torch.cuda.is_available() else torch.float32
	encoder = AutoModel.from_pretrained(
	HF_MODEL_NAME,
	torch_dtype=dtype,
	attn_implementation="sdpa",
	)

[gemini-code-assist]

The script contains a hardcoded absolute path pointing to a specific user's home directory (/home/aparnabg/...). This makes the script non-portable and unusable by other users or in other environments. Consider making the base directory configurable via an environment variable or command-line argument.

Suggested change

	SCRIPT_DIR="/home/aparnabg/orcd/scratch/all_project_files/action_sota_models/vjepa_crop"
	SCRIPT_DIR="${ACTION_SOTA_MODELS_DIR:-/home/aparnabg/orcd/scratch/all_project_files/action_sota_models/vjepa_crop}"